jorisvandenbossche commented on code in PR #34883:
URL: https://github.com/apache/arrow/pull/34883#discussion_r1157448809


##########
python/pyarrow/types.pxi:
##########
@@ -1494,6 +1494,135 @@ cdef class ExtensionType(BaseExtensionType):
         """
         return ExtensionScalar
 
+
+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:
+
+    >>> tensor_type = pa.FixedShapeTensorType(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 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 a fixed shape tensor extension type with permutation:
+    >>> tensor_type = pa.FixedShapeTensorType(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):

Review Comment:
   For all other types, we typically disallow calling the `__init__`, but 
provide a factory function to construct the type (eg `pa.timestamp(..)` instead 
of `pa.TimestampType(..)`). While that gives some extra API here, we should 
probably follow that pattern?



-- 
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]

Reply via email to