ENH: enable setitem dim2 test to work for EA with complex128 dtype

[MichaelTiemannOSC]

Feature Type

• Adding new functionality to pandas

• Changing existing functionality in pandas

• Removing existing functionality in pandas

Problem Description

I have read the threads related to DISCUSS/API: setitem-like operations should only update inplace (#39584) and friends (including #47577).

My problem arises from this test code from the Pint-Pandas test suite:

class TestSetitem(base.BaseSetitemTests):
    def test_setitem_2d_values(self, data):
        # GH50085                                                                                                                                                                                                                                                                                                                                                                                                       
        original = data.copy()
        df = pd.DataFrame({"a": data, "b": data})
	df.loc[[0, 1], :] = df.loc[[1, 0], :].values
        assert (df.loc[0, :] == original[1]).all()
        assert (df.loc[1, :] == original[0]).all()

As a result of PR #54441 I'm able to test Pint-Pandas with complex128 datatypes. PintArray EAs work perfectly with float and integer magnitudes, but fail in just this one case, with complex128. The failure starts in core/internals/managers.py in the fast_xs method:

        # GH#46406                                                                                                                                                                                                                                                                                                                                                                                                      
        immutable_ea = isinstance(dtype, ExtensionDtype) and dtype._is_immutable

        if isinstance(dtype, ExtensionDtype) and not immutable_ea:
            cls = dtype.construct_array_type()
            result = cls._empty((n,), dtype=dtype)

result becomes a PintArray backed by a FloatingArray :

<PintArray>
[<NA>, <NA>]
Length: 2, dtype: pint[nanometer]
(Pdb) result._data
<FloatingArray>
[<NA>, <NA>]
Length: 2, dtype: Float64

The FloatingArray comes when the PintArray initializer finds nothing helpful in either dtype nor values and falls back to creating a pd.array(values, ...):

    def __init__(self, values, dtype=None, copy=False):
        if dtype is None:
            if isinstance(values, _Quantity):
                dtype = values.units
            elif isinstance(values, PintArray):
                dtype = values._dtype
        if dtype is None:
            raise NotImplementedError

        if not isinstance(dtype, PintType):
            dtype = PintType(dtype)
        self._dtype = dtype

        if isinstance(values, _Quantity):
            values = values.to(dtype.units).magnitude
        elif isinstance(values, PintArray):
            values = values._data
        if isinstance(values, np.ndarray):
            dtype = values.dtype
            if dtype in dtypemap:
                dtype = dtypemap[dtype]
            values = pd.array(values, copy=copy, dtype=dtype)
            copy = False
        elif not isinstance(values, pd.core.arrays.numeric.NumericArray):
            values = pd.array(values, copy=copy)
        if copy:
            values = values.copy()
        self._data = values
        self._Q = self.dtype.ureg.Quantity

The fast_xs fails when result[rl] is not ready to accept the complex128 data coming from blk.iget((i, loc)):

        for blk in self.blocks:
            # Such assignment may incorrectly coerce NaT to None                                                                                                                                                                                                                                                                                                                                                        
            # result[blk.mgr_locs] = blk._slice((slice(None), loc))                                                                                                                                                                                                                                                                                                                                                     
            for i, rl in enumerate(blk.mgr_locs):
                result[rl] = blk.iget((i, loc))

As I see it, the problem is that we commit too soon to building our backing array with too-limited information.

@andrewgsavage
@topper-123
@jbrockmendel
@mroeschke

Feature Description

I will point out for the record:

• values = [v[loc] for v in self.arrays]

So we have everything we need within the environment of fast_xs. Should we use this knowledge as power to create a result that can hold slices of data beyond float64? Here's code that tries to use the fast path, but if an exception is raised, it does the sure thing:

        try:
            for blk in self.blocks:
                # Such assignment may incorrectly coerce NaT to None                                                                                                                                                                                                                                                                                                                                                    
                # result[blk.mgr_locs] = blk._slice((slice(None), loc))                                                                                                                                                                                                                                                                                                                                                 
                for i, rl in enumerate(blk.mgr_locs):
                    result[rl] = blk.iget((i, loc))
        except TypeError:
            if isinstance(dtype, ExtensionDtype) and not immutable_ea:
		values = [v[loc] for v in self.arrays]
                result = cls._from_sequence(values, dtype)
            else:
                raise TypeError

Alternative Solutions

I'm open to alternative solutions, but the above actually causes the test case to pass. Should I submit a PR?

Additional Context

No response

[jbrockmendel]

The FloatingArray comes when the PintArray initializer finds nothing helpful in either dtype nor values and falls back to creating a pd.array(values, ...):

From the wording it sounds like FloatingArray is not what you would actually want here? What type/dtype would you want here?

[MichaelTiemannOSC]

We have complex128 magnitudes to store, which cannot be stored in a FloatingArray. An array sized for complex128 would be great.

[MichaelTiemannOSC]

Apologies for slow responses...storms knocked out power yesterday and still no estimate of restoration. I did some tests and noticed that Pandas 2.0.2 passes this test with complex128, but not 2.1.0rc0. Still digging in...

[MichaelTiemannOSC]

Tracked it down: The fast_xs method in 2.1.0rc0 makes PintArray with [, ] whereas 2.0.x version makes it with [np.nan, np.nan]. The former demands FloatingArray, whereas the latter creates float64 array. 2.0.x silently converts (2+0j) to 2.0 before sticking it into the float64 array. 2.1.0rc0 properly complains that complex won't fit in float.

When using a "true" complex number (such as (2+1j) the 2.0.x version complains that float argument must be a string or a real number, not 'complex'.

When I get internet back I'll fix the test case. But the problem remains. One hack would be to not create an empty array, but an array holding the first value, then discard that value with [:1], preserving the underlying array type.

[lukemanley]

sorry if this is adding noise, but any chance #54508 fixes the fast_xs issue here?