Support for Compute Functions on Nested Arrays

[madhavajay]

Hi,
I have read through the docs and issues as best as I can and I am under the impression that its not possible to do compute functions on nested arrays.

I modified a group_by & aggregate example like so, putting the pa.array values into nested lists.

t = pa.table([
      pa.array(["a", "a", "b", "b", "c"]),
      pa.array([[1], [2], [3], [4], [5]]),
], names=["keys", "values"])

t.group_by("keys").aggregate([("values", "sum")])

The error is this:

ArrowNotImplementedError: Function 'hash_sum' has no kernel matching input types (array[list<item: int64>], array[uint32])

I assume this means the function doesn't know how to operate on a list? Is there a way to do this? I have large tensors which I can reshape into 1 dimension to store in a Record Batch, but I don't know how I can perform computations on their values. It seems like the other way is to use the Tensor type but it can't be used in a Record Batch or with compute can it?

The PyArrow zero copy from Numpy means this is an effective way to get data across the network using the IPC writer and its fairly easy to add other record types for custom meta data, but it would be a pity to have to then send this data back to numpy for all my computations and lose out on all that great SIMD parallelization.

Is there a better way?

Related links:
#4802
https://issues.apache.org/jira/browse/ARROW-1614

[pitrou]

The Tensor type is really a separate beast amongst Arrow containers. It cannot be used where arrays or record batches are expected. The Arrow compute functions currently all work on either arrays, record batches or scalars (and, of course, chunked arrays and tables).

You could open a JIRA for it but, to be honest, I doubt there's a lot of interest currently in improving the usefulness of tensors. The main use for them in PyArrow is in conjuction with https://sparse.pydata.org/en/stable/ or https://docs.scipy.org/doc/scipy/reference/sparse.html

cc @rok @mrkn for more information

[mrkn]

Supporting list types in hash_sum and supporting tensor types in a record batch are different issues.

@madhavajay Which feature can directly solve your problem?

[rok]

@madhavajay what kind of computation are you looking to do on the tensor values?
#8510 would introduce an easy way to go from Tensors to Arrays (making an array of tensors along one of the tensor dimensions) and back. However as Antoine says, there's currently no plans to add tensor operations in Arrow.

[pitrou]

hash_sum on lists would probably concatenate the lists, though? Though we would rather define another function name for this, I guess. @lidavidm

[lidavidm]

The other problem is that for non-commutative types, a hash_concat would not necessarily give meaningful results (and because of this we wouldn't want to have it in hash_sum)

I suppose with UDF support you could define tensor operations?

[pitrou]

Perhaps Arrow needs to add a Bag type that's like a List but unordered ;-)

[madhavajay]

@mrkn and @rok We are implementing a numpy like interface to large tensors with side car data to provide Differential Privacy and Secure Multi-Party Compute.

In the end we decided to only use the pyarrow Tensor format for zero copy, and then we store that in Capnp to send over the network because we need additional meta data and protobuf is max 2GB. We did try to shove it all in a single pyarrow record batch but that didn't really work. Finally for computation on the other end were using jax since its basically DL Tensor numpy data.

From a tabular perspective the data is essentially a row for each data subject (of whom we are protecting their privacy), like 1 large n-dim tensor, 2x similar ndim tensors (min and max providing bounds for the DP algorithms) but potentially with a lot of repeated data (so we made a custom datatype called lazyrepeatarray which removes duplicate dimensions), and finally there is information about the data subject, so in a perfect world that is all one single record but to make efficient use of the data and zero copy we need that all to be in a single array (column) for each data type rather than a record / row. However if there is no ability to do computation with pyarrow on that data then we just need to take it back out anyway. Currently were doing things like aggregate sum operations etc but we are implementing the entire suite of ops required for DL so we need numpy style flexibility.

Perhaps we were simply not utilising pyarrow correctly to best take advantage of what is possible. Also we are trying to avoid the serde which allows for code paths that support Python Objects aka pickle due to known security vulnerabilities.

I hope that helps.

[rok]

Thanks for the background info @madhavajay !

From a tabular perspective the data is essentially a row for each data subject (of whom we are protecting their privacy), like 1 large n-dim tensor, 2x similar ndim tensors (min and max providing bounds for the DP algorithms)

That sounds like a good fit for TensorArray as proposed in #8510 for c++ or as implemented in Python here

potentially with a lot of repeated data (so we made a custom datatype called lazyrepeatarray which removes duplicate dimensions)

Given you duplicate dimensions - would CSF sparse tensors help? Unlike CSR and CSC it will support n-dimensional tensors. Reduced size for transport and storage would be the main benefit as you'd still need dense forms for computation. Well some aggregate functions can be applied to sparse formats too..

However if there is no ability to do computation with pyarrow on that data then we just need to take it back out anyway. Currently were doing things like aggregate sum operations etc but we are implementing the entire suite of ops required for DL so we need numpy style flexibility.

Numpy does indeed seem the safest option. However as David mentioned there is ongoing work on Python UDFs, existing features are tested here.

[madhavajay]

@rok thanks for the information, it sounds like theres some great work going on here. I hope we can utilise more of this stuff in the future.

Regarding the Compressed Sparse Fibers SparseCSFTensor, firstly wow, thats really cool. For anyone arriving here who doesn't know what they are this blog is great: https://www.boristhebrave.com/2021/01/01/compressed-sparse-fibers-explained/

So now my only question is, while this seems like an optimal generalized solution for storage, how much computation is required to explode back out to the dense form in memory to do computation?

In our simple implementation since we are going by whole dimensions only, we can just use broadcast when necessary and then collapse back so the underlying data is just normal numpy arrays?

[rok]

So now my only question is, while this seems like an optimal generalized solution for storage, how much computation is required to explode back out to the dense form in memory to do computation?

I've not really benchmarked the conversion when implementing but I think it will heavily depend on your non-null distribution and even dimension order (!).
It should be pretty easy to benchmark though, just time:

sparse_tensor = pa.SparseCSFTensor.from_dense_numpy(np_array)
tensor = sparse_tensor.to_tensor()

In our simple implementation since we are going by whole dimensions only, we can just use broadcast when necessary and then collapse back so the underlying data is just normal numpy arrays?

I want to say yes, but I'm not 100% sure what you mean. Going from pa.Tensor to np.array and back should be zero copy AFAIK. Someone correct me if I'm wrong please!

[mrkn]

Going from pa.Tensor to np.array and back should be zero copy AFAIK.

It's true. Only metadata, that is the shape and the stride arrays are copied from Arrow's Tensor to NumPy's ndarray.

[madhavajay]

@rok and @mrkn thats awesome, I will get our team members to check this out!