jorisvandenbossche commented on code in PR #34883:
URL: https://github.com/apache/arrow/pull/34883#discussion_r1160507428
##########
python/pyarrow/array.pxi:
##########
@@ -3076,6 +3076,111 @@ cdef class ExtensionArray(Array):
return Array._to_pandas(self.storage, options, **kwargs)
+class FixedShapeTensorArray(ExtensionArray):
+ """
+ Concrete class for fixed shape tensor extension arrays.
+
+ Examples
+ --------
+ Define the extension type for tensor array
+
+ >>> import pyarrow as pa
+ >>> tensor_type = pa.fixed_shape_tensor(pa.int32(), [2, 2])
+
+ Create an extension array
+
+ >>> arr = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400]]
+ >>> storage = pa.array(arr, pa.list_(pa.int32(), 4))
+ >>> pa.ExtensionArray.from_storage(tensor_type, storage)
+ <pyarrow.lib.FixedShapeTensorArray object at ...>
+ [
+ [
+ 1,
+ 2,
+ 3,
+ 4
+ ],
+ [
+ 10,
+ 20,
+ 30,
+ 40
+ ],
+ [
+ 100,
+ 200,
+ 300,
+ 400
+ ]
+ ]
+ """
+
+ def to_numpy_ndarray(self):
+ """
+ Convert fixed shape tensor extension array to a numpy array (with
dim+1).
+ """
+ np_flat = np.asarray(self.storage.values)
+ numpy_tensor = np_flat.reshape((len(self),) + tuple(self.type.shape),
+ order='C')
Review Comment:
> Ah, actually I need to error if `permutation[0] != 0`.
No, the `self.type.permutation` is only that of the individual tensor
elements (given that you have a FixedShapeTensorArray, by definition the first
dimension of the n+1 dim ndarray is always the length of the array, and can't
be permuted)
> And if the permutation is for example `[0, 2, 1]` the user could still use
the information about the permutation to rearrange the ordering after the
reshape in `to_numpy_ndarray`.
Yes, and as mentioned before, I think we should certainly add examples how
the user can do this (if we decide to not automatically do it in
`to_numpy_ndarray`)
> Do we want to add a warning for dim_names not None also?
I wouldn't raise a warning for that, since it's a fact of life that numpy
arrays don't have dimension names, so that's just a consequence of calling this
method. We can mention that in the docstring, though, to be explicit.
--
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]
