Skip to content

GH-46771: [Python][C++] Implement pa.arange function to generate array sequences - #46778

Merged
raulcd merged 5 commits into
apache:mainfrom
raulcd:GH-46771
Jun 20, 2025
Merged

raulcd merged 5 commits into
apache:mainfrom
raulcd:GH-46771

Conversation

@raulcd

@raulcd raulcd commented Jun 11, 2025 •

Copy link
Copy Markdown
Member

Rationale for this change

When slicing arrays with non-trivial steps we were using numpy.arange to generate the indices for take. As numpy is an optional dependency, implementing it via Python caused a performance penalty. Creating a pyarrow function to build our own ranges that mimics Python range or numpy arange is useful for that uses case and might also be useful for other use cases. Currently we only generate Array[Int64] we could potentially generate more types.

What changes are included in this PR?

provide a pa.arange function that allows us to generate indices when slicing arrays.

Are these changes tested?

Yes new tests added.

Are there any user-facing changes?

No but a new pyarrow.arange function has been added.

@github-actions

This comment was marked as outdated.

@raulcd raulcd left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am unsure whether this is what you had in mind @pitrou but would appreciate some tips on possible improvements. I've been testing and it stills is slightly slower than the np.arange original implementation. With the original np.arange implementation:

In [2]: a = pa.array(np.arange(0, 2_000_000))

In [3]: %timeit a[slice(1, 1_000_000, 2)]
763 μs ± 4.28 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each

with the current PR one:

In [2]: a = pa.array(np.arange(0, 2_000_000))

In [3]: %timeit a[slice(1, 1_000_000, 2)]
1.08 ms ± 43.7 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

@raulcd raulcd changed the title [Python][C++] Implement pa.arange function to generate array sequences GH-46771: [Python][C++] Implement pa.arange function to generate array sequences Jun 11, 2025
@github-actions

Copy link
Copy Markdown

⚠️ GitHub issue #46771 has been automatically assigned in GitHub to PR creator.

@raulcd

raulcd commented Jun 11, 2025

Copy link
Copy Markdown
Member Author

It seems to be on-pair now:

In [2]: %timeit np.arange(1, 1_000_000, 2)
219 μs ± 10.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [3]: %timeit pa.arange(1, 1_000_000, 2)
212 μs ± 1.97 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Comment thread python/pyarrow/tests/test_array.py Resolved
Comment thread python/pyarrow/tests/test_array.py Resolved
Comment thread python/pyarrow/array.pxi Resolved
Comment thread python/pyarrow/array.pxi Resolved
@raulcd
raulcd marked this pull request as ready for review June 13, 2025 07:30
@raulcd
raulcd requested review from AlenkaF and rok as code owners June 13, 2025 07:30

@AlenkaF AlenkaF left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @raulcd for working on this!
It looks good on my end, but I'll let Antoine finish the review.

Comment thread python/pyarrow/array.pxi Resolved

@pitrou pitrou left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1, thanks for fixing this @raulcd

@conbench-apache-arrow

Copy link
Copy Markdown

After merging your PR, Conbench analyzed the 3 benchmarking runs that have been run so far on merge-commit b3e261d.

There were 122 benchmark results with an error:

There were no benchmark performance regressions. 🎉

The full Conbench report has more details.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants