ARROW-638: [C++] Complex Number Support via ExtensionTypes

[sjperkins]

xref https://issues.apache.org/jira/browse/ARROW-638

#10452 contains the PR implemented Complex Numbers as First Class Types

[github-actions]

https://issues.apache.org/jira/browse/ARROW-638

[pitrou]

Perhaps @jorisvandenbossche can answer the previous message, when he's back from vacation :-)

[jorisvandenbossche]

failing because there's no __arrow_ext_class__ attribute, which seems to be required by the Python layer

That might have been an oversight in the PR that introduced it (I don't think we currently have python tests using an extension type that is purely implemented in C++, and not using the (Py)ExtensionType subclassing mechanisms, although we should probably have such tests).
Can you check if moving the __arrow_ext_class__ base definition at

arrow/python/pyarrow/types.pxi

Lines 776 to 784 in c51e4a1

	def __arrow_ext_class__(self):
	"""Return an extension array class to be used for building or
	deserializing arrays with this extension type.
	This method should return a subclass of the ExtensionArray class. By
	default, if not specialized in the extension implementation, an
	extension type array will be a built-in ExtensionArray instance.
	"""
	return ExtensionArray

from the ExtensionType class to its parent class BaseExtensionType, fixes the problem?

[jmdeschenes]

Hello,

There is an issue with the approach:

on array.pxi

cdef class ExtensionArray(Array):
    """
    Concrete class for Arrow extension arrays.
    """


    @property
    def storage(self):
        cdef:
            CExtensionArray* ext_array = <CExtensionArray*>(self.ap)


        return pyarrow_wrap_array(ext_array.storage())
#
## LINES SKIPPED
#
    def to_numpy(self, **kwargs):
        """
        Convert extension array to a numpy ndarray.

        See Also
        --------
        Array.to_numpy
        """
        return self.storage.to_numpy(**kwargs)

on table.pxi

    def to_numpy(self):
        """
        Return a NumPy copy of this array (experimental).

        Returns
        -------
        array : numpy.ndarray
        """
        cdef:
            PyObject* out
            PandasOptions c_options
            object values


        if self.type.id == _Type_EXTENSION:
            storage_array = chunked_array(
                [chunk.storage for chunk in self.iterchunks()],
                type=self.type.storage_type

Both of these "strip" the Extension type sent to the CPP code. As such, the CPP code never knows that it is dealing with an extension.

If this is to be kept, fixed_size_list would need to convert into a proper 2D numpy array(That could have several benefits, it could be done only for primitive types at the start)

@jorisvandenbossche Do you think that is something that could be acceptable?

Otherwise, letting the CPP code handle the extension type could be another option.

@sjperkins Are you still working on this PR? Is there something I can help you with?

[sjperkins]

Thanks for checking in @jmdeschenes. I'm still working on it, but I haven't been able to give it attention recently. IIRC there are some issues I need to work through.

1. Most standard Numpy data types are First-Class Arrow types and the existing code handles conversion between them. Complex Number Arrow Types are Extension Types so their handling is a bit trickier and I haven't found a clean solution with the existing code, or with newer code that I've written yet.
2. I was very unsure about exposing C++ Extension Types in Python mostly because working in the CPython layer is slower and more opaque. It seems you've pointed out the reverse issue, in that Python Extension Types are stripped out during conversion to the C++ layer.

Any advice/help on the above 2 issues would be welcome. Note that I will likely only get another chance to look at this next year.

[jmdeschenes]

Thanks for the update.

With some changes, I was able to get it your branch to work and output a complex array. It is still preliminary. Let me know if you want to see it.

I think the best way forward is to let the CPP handle the extension types.

1. Don't strip the Extension information on the Cython side
2. CPP could handle future extension types each with their own potential Writer. Instead of stripping the Extension information on the cython side, let the CPP handle it(e.g. if the extension is known, use the Writer, if it doesn't, use the storage). This will require some rewrite on the CPP side as the extension type is used in the conversion code but I don't understand when it applies. To me this looks like the least intrusive option.

fixed_list_size_arrays generate a numpy array of pyarrow arrays. If fixed_list_size were to yield a 2D numpy array(which is directly mappable) there would be little needs for the extension side being handled at all. I guess this was done to ensure that the columns could be read into pandas. When calling to_numpy I would expect to get a 2D array.

As a separate issue, When trying to handle complex arrays from the python side only using extension arrays. It mostly works but the extension arrays don't play nice with ChunkedArrays

import pyarrow as pa
import pyarrow.parquet as pq
import numpy as np

class ComplexArray(pa.ExtensionArray):
    def to_numpy(self):
        array = self.storage.flatten().to_numpy(zero_copy_only=True).view('complex128')
        return array

class ComplexType(pa.ExtensionType):
    def __init__(self):
        pa.ExtensionType.__init__(self, pa.list_(pa.float64(), 2), "arrow.complex128")
    def __arrow_ext_serialize__(self):
        # since we don't have a parameterized type, we don't need extra
        # metadata to be deserialized
        return b''
    @classmethod
    def __arrow_ext_deserialize__(self, storage_type, serialized):
        # return an instance of this subclass given the serialized
        # metadata.
        return ComplexType()
    def __arrow_ext_class__(self):
        return ComplexArray

complex128 = ComplexType()
pa.register_extension_type(complex128)

array = np.array([1.0+1j, 2.0+2.0j, 3.0+3.0j]).view('float64')
float_array = pa.array(array, type=pa.float64())
list_array = pa.FixedSizeListArray.from_arrays(float_array, 2)
pa_array = ComplexArray.from_storage(complex128, list_array)
# This will yield the correct Numpy Array
assert (pa_array.to_numpy() == array).all()

It works but we get into issues when writing/reading from parquet

table = pa.Table.from_arrays([pa_array], names=['cpl'])
pq.write_table(table, "complex_test.parquet")
cpl_df = pq.read_table("complex_test.parquet")
# This line will create a numpy array of pyarrow arrays
cpl_df['cpl'].to_numpy()
# If you have a single chunk, this will work
assert (cpl_df['cpl'].chunks[0].to_numpy() == array).all()

# Otherwise the solution is more involved
start_length = 0
end_length = 0
max_length = cpl_df['cpl'].length()
final_numpy = np.zeros(shape=(max_length), dtype=np.complex128)
for chunk in cpl_df['cpl'].iterchunks():
    temp_numpy = chunk.to_numpy()
    end_length += temp_numpy.shape[0]
    final_numpy[start_length:end_length] = temp_numpy
    start_length = end_length
assert (final_numpy == array).all()

[sjperkins]

With some changes, I was able to get it your branch to work and output a complex array. It is still preliminary. Let me know if you want to see it.

Thanks, yes I would like to take a look.

I think the best way forward is to let the CPP handle the extension types.

1. Don't strip the Extension information on the Cython side

2. CPP could handle future extension types each with their own potential `Writer`. Instead of stripping the Extension information on the cython side, let the CPP handle it(e.g. if the extension is known, use the Writer, if it doesn't, use the storage). This will require some rewrite on the CPP side as the extension type is used in the conversion code but I don't understand when it applies. To me this looks like the least intrusive option.

Yes, I agree that this is the way forward. Mirroring what you've described back to you, this will necessitate API changes to the C++ Extension Framework to support Custom Writers?

fixed_list_size_arrays generate a numpy array of pyarrow arrays. If fixed_list_size were to yield a 2D numpy array(which is directly mappable) there would be little needs for the extension side being handled at all. I guess this was done to ensure that the columns could be read into pandas. When calling to_numpy I would expect to get a 2D array.

As a separate issue, When trying to handle complex arrays from the python side only using extension arrays. It mostly works but the extension arrays don't play nice with ChunkedArrays

import pyarrow as pa
import pyarrow.parquet as pq
import numpy as np

class ComplexArray(pa.ExtensionArray):
    def to_numpy(self):
        array = self.storage.flatten().to_numpy(zero_copy_only=True).view('complex128')
        return array

class ComplexType(pa.ExtensionType):
    def __init__(self):
        pa.ExtensionType.__init__(self, pa.list_(pa.float64(), 2), "arrow.complex128")
    def __arrow_ext_serialize__(self):
        # since we don't have a parameterized type, we don't need extra
        # metadata to be deserialized
        return b''
    @classmethod
    def __arrow_ext_deserialize__(self, storage_type, serialized):
        # return an instance of this subclass given the serialized
        # metadata.
        return ComplexType()
    def __arrow_ext_class__(self):
        return ComplexArray

complex128 = ComplexType()
pa.register_extension_type(complex128)

array = np.array([1.0+1j, 2.0+2.0j, 3.0+3.0j]).view('float64')
float_array = pa.array(array, type=pa.float64())
list_array = pa.FixedSizeListArray.from_arrays(float_array, 2)
pa_array = ComplexArray.from_storage(complex128, list_array)
# This will yield the correct Numpy Array
assert (pa_array.to_numpy() == array).all()

It works but we get into issues when writing/reading from parquet

table = pa.Table.from_arrays([pa_array], names=['cpl'])
pq.write_table(table, "complex_test.parquet")
cpl_df = pq.read_table("complex_test.parquet")
# This line will create a numpy array of pyarrow arrays
cpl_df['cpl'].to_numpy()
# If you have a single chunk, this will work
assert (cpl_df['cpl'].chunks[0].to_numpy() == array).all()

# Otherwise the solution is more involved
start_length = 0
end_length = 0
max_length = cpl_df['cpl'].length()
final_numpy = np.zeros(shape=(max_length), dtype=np.complex128)
for chunk in cpl_df['cpl'].iterchunks():
    temp_numpy = chunk.to_numpy()
    end_length += temp_numpy.shape[0]
    final_numpy[start_length:end_length] = temp_numpy
    start_length = end_length
assert (final_numpy == array).all()

It does work quite nicely in Python. In fact, I implemented something similar in a repository I'm working on: https://github.com/ska-sa/dask-ms/blob/55c987e5f00c24a82c16363f681a966e45590e06/daskms/experimental/arrow/extension_types.py#L145-L148. IIRC I ran into similar issues with chunked Extension Arrays.

Do you have an idea as to how this should be handled in the C++ layer?

[sjperkins]

With some changes, I was able to get it your branch to work and output a complex array. It is still preliminary. Let me know if you want to see it.

Hi @jmdeschenes, I'll be devoting some time to this PR over the next few weeks. Would it be possible to see your branch?

[pitrou]

For the record, we now have an official procedure for standardizing extension types. If you would like to resurrect this, you should post a proposal (or start discussing about such a proposal) on the ML.

[amol-]

Closing because it has been untouched for a while, in case it's still relevant feel free to reopen and move it forward 👍

[sjperkins]

This PR is very out of date, a new PR supporting this functionality will be required.