Add pa.fixedshapetensor factory function and update docstring examples

Author: AlenkaF

python/pyarrow/__init__.py

@@ -170,6 +170,7 @@ def print_entry(label, value):
                          union, sparse_union, dense_union,
                          dictionary,
                          run_end_encoded,
+                         fixedshapetensor,
                          field,
                          type_for_alias,
                          DataType, DictionaryType, StructType,

python/pyarrow/types.pxi

@@ -1499,78 +1499,23 @@ cdef class FixedShapeTensorType(BaseExtensionType):
     """
     Concrete class for fixed shape tensor extension type.
 
-    Parameters
-    ----------
-    value_type : DataType
-        Data type of individual tensor elements.
-    shape : tuple
-        The physical shape of the contained tensors.
-    dim_names : tuple
-        Explicit names to tensor dimensions.
-    permutation : tuple
-        Indices of the desired ordering of the original dimensions.
-
     Examples
     --------
-    >>> import pyarrow as pa
-
-    Create fixed shape tensor extension type:
+    Create an instance of fixed shape tensor extension type:
 
-    >>> tensor_type = pa.FixedShapeTensorType(pa.int32(), [2, 2])
-    >>> tensor_type
+    >>> import pyarrow as pa
+    >>> pa.fixedshapetensor(pa.int32(), [2, 2])
     FixedShapeTensorType(extension<arrow.fixed_shape_tensor>)
 
-    Inspect the data type:
-
-    >>> tensor_type.value_type
-    DataType(int32)
-    >>> tensor_type.shape
-    [2, 2]
-
-    Create a fixed shape tensor extension type with names of tensor dimensions:
-
-    >>> tensor_type = pa.FixedShapeTensorType(pa.int8(), (2, 2, 3), dim_names=['C', 'H', 'W'])
-    >>> tensor_type.dim_names
-    [b'C', b'H', b'W']
+    Create an instance of fixed shape tensor extension type with
+    permutation:
 
-    Create a fixed shape tensor extension type with permutation:
-    >>> tensor_type = pa.FixedShapeTensorType(pa.int8(), (2, 2, 3), permutation=[0, 2, 1])
+    >>> tensor_type = pa.fixedshapetensor(pa.int8(), (2, 2, 3),
+    ...                                   permutation=[0, 2, 1])
     >>> tensor_type.permutation
     [0, 2, 1]
     """
 
-    def __init__(self, DataType value_type, shape, dim_names=None, permutation=None):
-        """
-        Initialize an fixed shape tensor extension type instance.
-
-        This should be called at the end of the subclass'
-        ``__init__`` method.
-        """
-        cdef:
-            vector[int64_t] c_shape
-            vector[int64_t] c_permutation
-            vector[c_string] c_dim_names
-            shared_ptr[CDataType] tensor_ext_type
-
-        assert value_type is not None
-        assert shape is not None
-
-        for i in shape:
-            c_shape.push_back(i)
-
-        if permutation is not None:
-            for i in permutation:
-                c_permutation.push_back(i)
-
-        if dim_names is not None:
-            for x in dim_names:
-                c_dim_names.push_back(tobytes(x))
-
-        tensor_ext_type = GetResultValue(CFixedShapeTensorType.Make(
-            value_type.sp_type, c_shape, c_permutation, c_dim_names))
-
-        self.init(tensor_ext_type)
-
     cdef void init(self, const shared_ptr[CDataType]& type) except *:
         BaseExtensionType.init(self, type)
         self.tensor_ext_type = <const CFixedShapeTensorType*> type.get()
@@ -1607,17 +1552,15 @@ cdef class FixedShapeTensorType(BaseExtensionType):
         """
         Serialized representation of metadata to reconstruct the type object.
         """
-        metadata = self.tensor_ext_type.Serialize()
-        return metadata
+        return self.tensor_ext_type.Serialize()
 
     @classmethod
     def __arrow_ext_deserialize__(self, storage_type, serialized):
         """
         Return an FixedShapeTensor type instance from the storage type and serialized
         metadata.
         """
-        tensor_ext_type = self.tensor_ext_type.Deserialize(storage_type, serialized)
-        return tensor_ext_type
+        return self.tensor_ext_type.Deserialize(storage_type, serialized)
 
     def __arrow_ext_class__(self):
         return FixedShapeTensorArray
@@ -4672,6 +4615,100 @@ def run_end_encoded(run_end_type, value_type):
     return pyarrow_wrap_data_type(ree_type)
 
 
+def fixedshapetensor(DataType value_type, shape, dim_names=None, permutation=None):
+    """
+    Create instance of fixed shape tensor extension type with shape and optional
+    names of tensor dimensions and indices of the desired ordering.
+
+    Parameters
+    ----------
+    value_type : DataType
+        Data type of individual tensor elements.
+    shape : tuple
+        The physical shape of the contained tensors.
+    dim_names : tuple, default None
+        Explicit names to tensor dimensions.
+    permutation : tuple, default None
+        Indices of the desired ordering of the original dimensions.
+
+    Examples
+    --------
+    Create an instance of fixed shape tensor extension type:
+
+    >>> import pyarrow as pa
+    >>> tensor_type = pa.fixedshapetensor(pa.int32(), [2, 2])
+    >>> tensor_type
+    FixedShapeTensorType(extension<arrow.fixed_shape_tensor>)
+
+    Inspect the data type:
+
+    >>> tensor_type.value_type
+    DataType(int32)
+    >>> tensor_type.shape
+    [2, 2]
+
+    Create a table with fixed shape tensor extension array:
+
+    >>> arr = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400]]
+    >>> storage = pa.array(arr, pa.list_(pa.int32(), 4))
+    >>> tensor = pa.ExtensionArray.from_storage(tensor_type, storage)
+    >>> pa.table([tensor], names=["tensor_array"])
+    pyarrow.Table
+    tensor_array: extension<arrow.fixed_shape_tensor>
+    ----
+    tensor_array: [[[1,2,3,4],[10,20,30,40],[100,200,300,400]]]
+
+    Create an instance of fixed shape tensor extension type with names
+    of tensor dimensions:
+
+    >>> tensor_type = pa.fixedshapetensor(pa.int8(), (2, 2, 3),
+    ...                                   dim_names=['C', 'H', 'W'])
+    >>> tensor_type.dim_names
+    [b'C', b'H', b'W']
+
+    Create an instance of fixed shape tensor extension type with
+    permutation:
+
+    >>> tensor_type = pa.fixedshapetensor(pa.int8(), (2, 2, 3),
+    ...                                   permutation=[0, 2, 1])
+    >>> tensor_type.permutation
+    [0, 2, 1]
+
+    Returns
+    -------
+    type : FixedShapeTensorType
+    """
+
+    cdef:
+        vector[int64_t] c_shape
+        vector[int64_t] c_permutation
+        vector[c_string] c_dim_names
+        shared_ptr[CDataType] c_tensor_ext_type
+
+    assert value_type is not None
+    assert shape is not None
+
+    for i in shape:
+        c_shape.push_back(i)
+
+    if permutation is not None:
+        for i in permutation:
+            c_permutation.push_back(i)
+
+    if dim_names is not None:
+        for x in dim_names:
+            c_dim_names.push_back(tobytes(x))
+
+    cdef FixedShapeTensorType out = FixedShapeTensorType.__new__(FixedShapeTensorType)
+
+    c_tensor_ext_type = GetResultValue(CFixedShapeTensorType.Make(
+        value_type.sp_type, c_shape, c_permutation, c_dim_names))
+
+    out.init(c_tensor_ext_type)
+
+    return out
+
+
 cdef dict _type_aliases = {
     'null': null,
     'bool': bool_,