BUG: 1D slices over extension types turn into N-dimensional slices over ExtensionArrays · Issue #42430 · pandas-dev/pandas (original) (raw)

Note: Please read this guide detailing how to provide the necessary information for us to reproduce your bug.

Code Sample, a copy-pastable example

class LoggingStringArray(pd.arrays.StringArray): """Extend StringArray to log arguments to getitem""" def getitem(self, item): print(f"Argument to getitem is '{item}'") return super().getitem(item)

df = pd.DataFrame({"col1": LoggingStringArray(np.array(["hello", "world"], dtype=object))}) _ = df.iloc[:1]

Expected output:

Argument to __getitem__ is 'slice(None, 1, None)'

Actual output:

Argument to __getitem__ is '(Ellipsis, slice(None, 1, None))'

Problem description

If a DataFrame contains extension types, the block manager turns a 1D slice in the form DataFrame.iloc[slice(x, y, z)] into an N-dimensional slice over the ExtensionArray of the form __getitem__((Ellipsis, slice(x, y, z)) . This problem breaks the __getitem__() method on one of our extension types in Text Extensions for Pandas.

The __getitem__ methods in Pandas' built-in extension types happens to work when fed this kind of n-dimensional slice
--- they just ignore the Ellipsis in the first position of the tuple. But I don't think they're supposed to accept that type of input.

The root cause is the new method Block.getitem_block_index(), added during some performance optimizations introduced in #40353:

def getitem_block_index(self, slicer: slice) -> Block:
    """
    Perform __getitem__-like specialized to slicing along index.

    Assumes self.ndim == 2
    """
    # error: Invalid index type "Tuple[ellipsis, slice]" for
    # "Union[ndarray, ExtensionArray]"; expected type "Union[int, slice, ndarray]"
    new_values = self.values[..., slicer]  # type: ignore[index]
    return type(self)(new_values, self._mgr_locs, ndim=self.ndim)

This method assumes that the contents of the block are 2-dimensional. The ExtensionBlock class inherits this method from EABackedBlock, which inherits it from Block.

Output of pd.show_versions()

INSTALLED VERSIONS

commit : f00ed8f
python : 3.7.10.final.0
python-bits : 64
OS : Darwin
OS-release : 20.5.0
Version : Darwin Kernel Version 20.5.0: Sat May 8 05:10:33 PDT 2021; root:xnu-7195.121.3~9/RELEASE_X86_64
machine : x86_64
processor : i386
byteorder : little
LC_ALL : None
LANG : en_US.UTF-8
LOCALE : en_US.UTF-8

pandas : 1.3.0
numpy : 1.21.0
pytz : 2021.1
dateutil : 2.8.1
pip : 21.1.3
setuptools : 52.0.0.post20210125
Cython : None
pytest : 6.2.4
hypothesis : None
sphinx : 4.0.2
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : None
psycopg2 : None
jinja2 : 3.0.1
IPython : 7.25.0
pandas_datareader: None
bs4 : None
bottleneck : None
fsspec : None
fastparquet : None
gcsfs : None
matplotlib : 3.4.2
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : 4.0.1
pyxlsb : None
s3fs : None
scipy : 1.7.0
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
xlwt : None
numba : None