rok commented on code in PR #38008:
URL: https://github.com/apache/arrow/pull/38008#discussion_r1410029983
##########
python/pyarrow/tests/test_extension_type.py:
##########
@@ -1472,10 +1618,24 @@ def test_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):
+def test_tensor_type_str(tensor_type, text, pickle_module):
Review Comment:
Done.
##########
python/pyarrow/types.pxi:
##########
@@ -4948,6 +5056,121 @@ def fixed_shape_tensor(DataType value_type, shape,
dim_names=None, permutation=N
return out
+def variable_shape_tensor(DataType value_type, ndim, dim_names=None,
permutation=None,
+ uniform_shape=None):
+ """
+ Create instance of variable shape tensor extension type with number of
+ dimensions and optional names of tensor dimensions and indices of the
+ desired logical ordering of dimensions.
+
+ Parameters
+ ----------
+ value_type : DataType
+ Data type of individual tensor elements.
+ ndim : integer
+ The number of dimensions of the contained tensors.
+ dim_names : tuple or list of strings, default None
+ Explicit names to tensor dimensions.
+ permutation : tuple or list integers, default None
+ Indices of the desired ordering of the original dimensions.
+ The indices contain a permutation of the values ``[0, 1, .., N-1]``
where
+ N is the number of dimensions. The permutation indicates which
dimension
+ of the logical layout corresponds to which dimension of the physical
tensor.
+ For more information on this parameter see
+ :ref:`fixed_shape_tensor_extension`.
+ uniform_shape : tuple or list of integers, default None
+ Shape of dimensions that are guaranteed to stay constant over all
tensors
+ in the array if all their non-dimensions sizes were replaced by None.
+
+ Examples
+ --------
+ Create an instance of variable shape tensor extension type:
+
+ >>> import pyarrow as pa
+ >>> tensor_type = pa.variable_shape_tensor(pa.int32(), 2)
+ >>> tensor_type
+ VariableShapeTensorType(extension<arrow.variable_shape_tensor>)
+
+ Inspect the data type:
+
+ >>> tensor_type.value_type
+ DataType(int32)
+ >>> tensor_type.ndim
+ 2
+
+ Create a table with variable shape tensor extension array:
+
+ >>> fields = [pa.field("shape", pa.list_(pa.uint32(), 2)),
pa.field("data", pa.list_(pa.int32()))]
+ >>> storage = pa.array([([2, 3], [1, 2, 3, 4, 5, 6]), ([1, 2], [7, 8])],
type=pa.struct(fields))
+ >>> tensor = pa.ExtensionArray.from_storage(tensor_type, storage)
+ >>> pa.table([tensor], names=["tensor_array"])
+ pyarrow.Table
+ tensor_array: extension<arrow.variable_shape_tensor>
+ ----
+ tensor_array: [ -- is_valid: all not null
+ -- child 0 type: fixed_size_list<item: uint32>[2]
+ [[2,3],[1,2]]
+ -- child 1 type: list<item: int32>
+ [[1,2,3,4,5,6],[7,8]]]
+
+ Create an instance of variable shape tensor extension type with names
+ of tensor dimensions:
+
+ >>> tensor_type = pa.variable_shape_tensor(pa.int8(), 3,
+ ... dim_names=['C', 'H', 'W'])
+ >>> tensor_type.dim_names
+ ['C', 'H', 'W']
+
+ Create an instance of variable shape tensor extension type with
+ permutation:
+
+ >>> tensor_type = pa.variable_shape_tensor(pa.int8(), 3,
+ ... permutation=[0, 2, 1])
+ >>> tensor_type.permutation
+ [0, 2, 1]
+
+ Returns
+ -------
+ type : VariableShapeTensorType
+ """
+
+ cdef:
+ uint32_t c_ndim
+ vector[int64_t] c_permutation
+ vector[c_string] c_dim_names
+ 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
+
+ c_ndim = ndim
+
+ 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))
+
+ if uniform_shape is not None:
+ for x in uniform_shape:
+ if x is None:
+ c_uniform_shape.push_back(<optional[int64_t]>nullopt)
+ else:
+ c_uniform_shape.push_back(<optional[int64_t]>(<int64_t>x))
+
+ cdef VariableShapeTensorType out =
VariableShapeTensorType.__new__(VariableShapeTensorType)
+
+ c_tensor_ext_type = GetResultValue(CVariableShapeTensorType.Make(
+ value_type.sp_type, c_ndim, c_permutation, c_dim_names,
c_uniform_shape))
Review Comment:
Done.
--
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]
