GH-32007: [Python] Support arithmetic on arrays and scalars

[rmnskb]

Rationale for this change

Please see #32007, currently, neither arrays nor scalars support Python-native arithmetic operations, such as array + array, it has to be done via pyarrow.compute API. This PR strives to fix this with custom dunder methods.

What changes are included in this PR?

Implemented dunder methods

Are these changes tested?

Yes

Are there any user-facing changes?

Possibility to use Python operators directly instead of calling the pyarrow.compute API.

• GitHub Issue: [Python] Support arithmetic on arrays and scalars #32007

[rmnskb]

@pitrou @AlenkaF @rok
Sorry for pinging like this, I wanted to understand the scope correctly. So far, I went with the most straightforward implementation and added dunders for all operations that were listed in the compute docs and had respective dunders available. Do you think this implementation is enough or should I go for other dunders as well, e.g bitwise operations? After reviewing the original discussion, I've seen @jorisvandenbossche suggesting checked versions of the functions, I think I will go with them as well.
Regarding the documentation:

1. Should I add docstrings to dunder methods? Given that user won't see them in their IDE, for example.
2. Should the documentation for arrays/scalars be updated to reflect that the native operators are now supported?
   Thanks for your input!

[rok]

Sorry for late reply @rmnskb. I think all dunders that semantically map to our existing kernels are fair game! I'm not sure what these would be, but the set you have now looks decently sized.

Should I add docstrings to dunder methods? Given that user won't see them in their IDE, for example.

We might be able to add docstrings via type annotations. Or perhaps it's possible to copy them at runtime from the respective kernels? (unsure that is possible). In any case we probably don't want to duplicate docstrings in general.

Should the documentation for arrays/scalars be updated to reflect that the native operators are now supported?

Yes, that would be good to have!

[rmnskb]

Sorry for late reply @rmnskb. I think all dunders that semantically map to our existing kernels are fair game! I'm not sure what these would be, but the set you have now looks decently sized.

Ok, I was also thinking about covering as many dunder methods as possible, but did not want to go out of scope for this issue.

We might be able to add docstrings via type annotations. Or perhaps it's possible to copy them at runtime from the respective kernels? (unsure that is possible). In any case we probably don't want to duplicate docstrings in general.

Copying the docstrings sounds like a good idea, I will look further into that.

[AlenkaF]

Thank you for working on this @rmnskb and sorry for replying late!
Pinging us when you have a question is totally fine and in fact necessary ;)

I think the current scope is solid. I am thinking of basic comparisons and the discussion in the issue around __eq__. That would also be good to tackle but maybe as a separate PR?

Yes, it would be important to update the documentation (User Guide which is the part that is not API reference related).

One comment after looking at the code would be to add tests that cover various type combinations (array, scalar, Python types, unsupported types ...)

[rok]

I think the current scope is solid. I am thinking of basic comparisons and the discussion in the issue around __eq__. That would also be good to tackle but maybe as a separate PR?

I think we already have __eq__: https://github.com/apache/arrow/pull/7737/files.

One comment after looking at the code would be to add tests that cover various type combinations (array, scalar, Python types, unsupported types ...)

That's a good point, __eq__ currently raises if the types of compared don't match. Good test coverage would also help find edge cases.

    def __eq__(self, other):
        try:
            return self.equals(other)
        except TypeError:
            # This also handles comparing with None
            # as Array.equals(None) raises a TypeError.
            return NotImplemented

[AlenkaF]

I think we already have __eq__: https://github.com/apache/arrow/pull/7737/files.

Yes, true. I was under the impression that we might want to change the use of equals() method with the equal compute function? From the issue:

I think if we add those, we should also change the behaviour of __eq__, although that is something that will require a long deprecation cycle.

[rok]

Oh, I managed to forget this and thought array.equals dispatches to compute.equals already 🤦.

It seems like a good opportunity to change this! :)

[rmnskb]

We might be able to add docstrings via type annotations. Or perhaps it's possible to copy them at runtime from the respective kernels? (unsure that is possible). In any case we probably don't want to duplicate docstrings in general.

I've looked more into this, and based on these two discussions, we'd have to copy them directly from the underlying functions at runtime via importing the pyarrow.compute module, which AFAIK we don't have access to, since we call them with _pc().call_function().
I've added the docstrings to respective classes which notes that the API now support calling the Python-native operators directly, as well as provides some examples of that. Please let me know if that's sufficient.

[rmnskb]

I think the current scope is solid. I am thinking of basic comparisons and the discussion in the issue around __eq__. That would also be good to tackle but maybe as a separate PR?

I'll be happy to work on it, once I'm done with this implementation :)

Yes, it would be important to update the documentation (User Guide which is the part that is not API reference related).

Are you talking about this file? Or is there somewhere else, where I can put this information?

One comment after looking at the code would be to add tests that cover various type combinations (array, scalar, Python types, unsupported types ...)

Yes, I agree. Are there any particular unsupported types I should include that you have in mind?

[AlenkaF]

I'll be happy to work on it, once I'm done with this implementation :)

Thank you!

Are you talking about this file? Or is there somewhere else, where I can put this information?

I was thinking more about the compute page in our User Guide: https://github.com/apache/arrow/blob/main/docs/source/python/compute.rst

Are there any particular unsupported types I should include that you have in mind?

No, not really. What comes to mind are strings and/or nested for some arithmetic functions.

[rmnskb]

Are you talking about this file? Or is there somewhere else, where I can put this information?

I was thinking more about the compute page in our User Guide: https://github.com/apache/arrow/blob/main/docs/source/python/compute.rst

I mentioned the newly implemented operators in the documentation. Please let me know if it makes sense. I'm not sure whether we should list all the implemented methods, on the other hand, I don't want to leave users guessing what exactly can they use.

[rok]

One comment after looking at the code would be to add tests that cover various type combinations (array, scalar, Python types, unsupported types ...)

Yes, I agree. Are there any particular unsupported types I should include that you have in mind?

An integer based extension type maybe.

[rmnskb]

@AlenkaF @rok thanks for your input on this PR, I've marked it as ready, so I think it's up for a review, please let me know if there is anything else I should include, thanks! :)

[AlenkaF]

Great work, thank you @rmnskb!
I have one comment regarding the documentation and one minor nit. Also, could we add one more test that includes mixed type arithmetics (like in docstring examples: array + scalar, etc.)?

[rmnskb]

Hey, @AlenkaF, good points! Thanks for the review. I've added another test and updated the docs as per your suggestion.

[AlenkaF]

I think this is a great improvement, thank you for working on it @rmnskb !
Will wait for another review before merging. @rok, would you have time to have a look?

[AlenkaF]

@pitrou @raulcd looking for another review before merging, if anybody has time.

[jorisvandenbossche]

Two general quick comments:

• The & / | operators are currently mapped to the bitwise kernels, but in array-land, I think the most common use case for those is actually for logical (boolean) operations. So as a user, I would expect those to also work on bool type.
• Do we want to add the round/floor/ceil/trunc rounding dunders? Numpy does not support those, so here I wouldn't have expected them. But maybe it can also be useful (I think it find it a bit strange looking that a stdlib builtin function returns a pyarrow array, but that is probably a question of getting used to it). One difference is that the builtin is documented to return an int, and our kernels don't do that (float stays float)

[rmnskb]

The & / | operators are currently mapped to the bitwise kernels, but in array-land, I think the most common use case for those is actually for logical (boolean) operations. So as a user, I would expect those to also work on bool type.

That is true, I also though about it, but it seems to me that being consistent with other dunders also makes sense. Should we probably document this? So the users know what to expect.

Do we want to add the round/floor/ceil/trunc rounding dunders? Numpy does not support those, so here I wouldn't have expected them. But maybe it can also be useful (I think it find it a bit strange looking that a stdlib builtin function returns a pyarrow array, but that is probably a question of getting used to it). One difference is that the builtin is documented to return an int, and our kernels don't do that (float stays float)

Fair point. I think documenting the expected behavior might be a good idea given that Python is quite forgiving when it comes to data types. What do you think? The other option is to explicitly convert the output of this dunders to integers to stay consistent with the builtin docs.

[pitrou]

Do we want to gold-plate this? We can keep the bitwise operators, remove the round/floor/etc. dunders, and perhaps improve things later if people request it.

[AlenkaF]

We can split this PR in two and move round/floor/etc., ... in a follow-up one if there will be requests for them.

[AlenkaF]

Thank you for the update @rmnskb!
The documentation CI failures do not look related. Sphinx job seems to start failing with pandas 3.0 and needs a separate fix (will open an issue if it is not opened yet). The Debian job is connected to R documentation, I need to investigate that also (as a separate issue).

I think the latest change addresses the last comments, @pitrou or @raulcd - mind giving another look at the changes?

[pitrou]

If we really want to assert the "checked" aspect, then we should also include a case where overflow occurs.

[rmnskb]

The overall intention of this was to test that the dunder methods have the same output given the same input (kernels as well). Albeit, your suggestion is valid, so added extra tests to cover whether dunder methods indeed overflow.

[pitrou]

Can I simply call arr1 - 42 or would that not work?

[rmnskb]

It would work, yes. I wanted to show in the docstrings that the users can also use explicit scalars, that's why I went with this.

[pitrou]

+1, thank you! @AlenkaF Do you want to give this the final go?