Pyarrow compute against extension arrays

[Spencer Nelson <sw...@uw.edu>]

I'm trying to work with new FixedShapeTensors from Pyarrow. I'd like to use
tensors to store covariance matrices, so I'd like to get the diagonal
elements of my (6x6) matrix.

I'm working through the most naive version of this, which is to hardcode
the selections of the (0,0), (1, 1), ... (6, 6) elements. I do that
with "pc.list_element(pc.list_element(my_array,
0), 0)".

When I try that, I get an error that I think might expose a weakness in the
interactions between pyarrow.compute and extension types:

pyarrow.lib.ArrowNotImplementedError: Function 'list_element' has no kernel
matching input types (extension<arrow.fixed_shape_tensor>, int64)

It's no surprise that there is no matching kernel for an extension
type, although I'd sort of expect that pyarrow would inspect the *storage* of
the extension array type.

Is there a better way for me to approach this?