REGR: Performance regression in axis=1 DataFrame ops

[rhshadrach]

Ref: #51335 (comment)

• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

This is a work in progress; if this looks like a viable way forward it at least needs tests (and maybe fixes) for non-numeric dtypes, a whatsnew, and some cleanups. This PR gets around the performance regression by operating on NumPy dtypes only in cases where we can be certain the op can correctly be computed (namely NumPy and masked EA dtypes).

On 1.5.x, the axis=1 would result in NumPy float dtype, even when e.g. summing all integers. Here the result remains a masked EA and ops like sum, prod don't upcast to float.

cc @jbrockmendel @jorisvandenbossche

[jbrockmendel]

can this be len(df.columns)

[jbrockmendel]

also we only need to care about common_dtype in axis=1 case right? can we avoid these checks otherwise?

[jbrockmendel]

Yikes. I'm leaning towards just acknowledging that these ops with MaskedDtype will always perform poorly.

[rhshadrach]

Yikes.

Given the current state of this PR I don't disagree, but it can be somewhat cleaned up. I'm curious if it's the logic after the NumPy computation (L10538-10557) that gives you that reaction, as I don't see anyway around that part.

I'm thinking this is worth it as it fixes not only performance but also some dtype issues:

values = pd.array([True, pd.NA], dtype="boolean")
df = DataFrame({"a": values, "b": values})

# 2.0.x
print(df.min(axis=1))
# 0    True
# 1    <NA>
# dtype: object

# PR
print(df.min(axis=1))
0    True
1    <NA>
dtype: boolean

That said, if this doesn't look viable for the masked dtypes, perhaps some thought could go into changing the EA API to better support axis=1 ops (even for third party EAs)?

[jbrockmendel]

not really relevant here but in general df._values is better than df.values, preserves e.g. DatetimeArray[tzaware]

[jbrockmendel]

I'm thinking this is worth it as it fixes not only performance but also some dtype issues:

Thanks for pointing this out, I was working under the wrong impression. For now I'm going to ignore the perf issue (as mentioned, I'm basically willing to write it off) and only focus on the bugs (which I expect includes #42895).

IIUC the crux of the problem is in Block.reduce where we do

        if self.values.ndim == 1:
            # TODO(EA2D): special case not needed with 2D EAs
            res_values = np.array([[result]])
        else:
            res_values = result.reshape(-1, 1)

In the self.values.ndim == 1 case result is a scalar that has lost the relevant dtype information.

ndarray reductions have a keepdims keyword we could adapt so that a 1D EA reduces to a length-1 1D EA instead of a scalar. I think something like that would solve the General Case problem here.

Shorter term what if we tried a kludgy patch in Block.reduce checking for MaskedEADtype?

[Obligatory mention that 2D EAs would have prevented this and so so many other issues.]

[jbrockmendel]

The df.min case looks much improved by

        if self.values.ndim == 1:
            # TODO(EA2D): special case not needed with 2D EAs
            if isinstance(self.dtype, BaseMaskedDtype):
                # TODO: use _maybe_mask_result?
                if self.dtype.kind == "b":
                    res_values = type(self.values)._from_sequence([result])
                else:
                    res_values = type(self.values)._from_sequence([result])
            else:
                res_values = np.array([[result]])
        else:
            res_values = result.reshape(-1, 1)

though i think that only works for a subset of cases

[github-actions]

This pull request is stale because it has been open for thirty days with no activity. Please update and respond to this comment if you're still interested in working on this.