Astropy CSV table reader using pyarrow

[taldcroft]

Description

This pull request is an implementation of a fast CSV reader for astropy that uses pyarrow.csv.read_csv. This was discussed in #16869. In particular, a speed-up of around 10x over pandas and the astropy fast reader is noted in this profiling by @dhomeier.

The goal for this reader is to make an interface that will be familiar to astropy io.ascii users, while exposing some additional features brought by pyarrow read_csv. The idea is to keep the interface clean and consistent with astropy.

A quick demonstration notebook that you can use to play with this is at: https://gist.github.com/taldcroft/ac15bc516a7bf7c76f9eec644c787298

Fixes #16869

Related

pandas-dev/pandas#54466

Please DO squash and merge, the individual commits are not valuable here.

• By checking this box, the PR author has requested that maintainers do NOT use the "Squash and Merge" button. Maintainers should respect this when possible; however, the final decision is at the discretion of the maintainer that merges the PR.

[github-actions]

Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

• Do the proposed changes actually accomplish desired goals?
• Do the proposed changes follow the Astropy coding guidelines?
• Are tests added/updated as required? If so, do they follow the Astropy testing guidelines?
• Are docs added/updated as required? If so, do they follow the Astropy documentation guidelines?
• Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see instructions for rebase and squash.
• Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the bot.
• Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
• Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
• At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) before merge.

[github-actions]

👋 Thank you for your draft pull request! Do you know that you can use [ci skip] or [skip ci] in your commit messages to skip running continuous integration tests until you are ready?

[pllim]

Thanks! I want to benchmark this but does that mean we need to install pyarrow in https://github.com/astropy/astropy/blob/main/.github/workflows/ci_benchmark.yml ?

[taldcroft]

Thanks! I want to benchmark this but does that mean we need to install pyarrow in https://github.com/astropy/astropy/blob/main/.github/workflows/ci_benchmark.yml ?

There are one-time benchmarks here: #16869 (comment). These demonstrate that pyarrow read_csv() appears to be a factor of 10 faster than any other readers.

[mhvk]

Nice! I like the general idea; my only more major comment is that I'm not sure in this initial stage one should add the commented-line skipper.

For follow-up, I guess, would be to make this the default "first try" if pyarrow is available, and then deprecate the fast reader?

It does seem Table.{from,to}_pyarrow methods would be reasonable, but better as follow-up.

[hamogu]

Since @dhomeier has shown pyarrow to be significantly faster, it would be good to have it for the biggest tables. And this is a relatively thin wrapper just to match the API we are used to, so why not?
I do wonder (similar to @mhvk ) how far it makes sense to go to have capabilities that are not native to pyarrow (e.g. comment characters). Is it worth the pure-Python preprocessing at all? Would that dilute the advertised point "this is super fast and super-memory efficient, so use it for tables in the GB range"?

For smaller tables we have other established solutions which are more flexible (not the least our own pure-python readers and our own C reader). How many GB-sized tables are there in the wild with commented lines that are not in the header? I'm just worried about user confusion along the lines of "It's reading this table just fine and that table that's almost identical (but with comment lines) crashes with a Python out-of-memory error". Of course, that only applies to the biggest tables of them all. For csv files in the 0.5-1GB rage, this is probably still be faster AND would fit into memory (and maybe not be too slow) on modern machines. So it's a trade-off.

[neutrinoceros]

This may be a bit too technical for a first round of review but I wanted to get this kind of feedback in early too so it doesn't grow into too much of a pain later:
Here are a couple suggestions and comments mostly about type annotations and internal consistency.

[taldcroft]

Thanks for the great comments! I think I've addressed them all, or at least responded. Sounds like I have agreement to keep going ahead on this and start working on tests, docs etc?

[taldcroft]

@mhvk @neutrinoceros - I've addressed all the first round of comments.

I lost track of the feature freeze date until the recent announcement, but I hope this can get into this release. I'm planning to add testing and documentation on Monday so I would be very grateful if you are able to make time to look at this again early this week.

[taldcroft]

One unfortunate thing I just noticed is that my current profiling shows that this full implementation is only a factor of 2x faster than the fast ASCII reader.

dat = simple_table(size=100000, cols=20)
dat.write("junk.csv", format="ascii.basic", overwrite=True, delimiter=",")

%timeit Table.read("junk.csv", format="ascii.csv", fast_reader="force", delimiter=",", guess=False)
103 ms

%timeit Table.read("junk.csv", format="pyarrow.csv")
51 ms

It turns out that converting an object arrays of strings into a numpy string array is the bottleneck.

[mhvk]

It turns out that converting an object arrays of strings into a numpy string array is the bottleneck.

This may well go away eventually, with numpy's new StringDType - there's talk about trying to have a copy-less version with pyarrow. Anyway, a factor 2 is still nice! And presumably a lot better for mostly numeric data?

[taldcroft]

Profiling the new implementation pyarrow.csv vs fast io.ascii gives:

• About 30% faster for all string columns
• About 16x faster for all int and float columns

For a table with two commented lines at the top of the file:

• About 5x faster for all int, float columns, with no change in memory use.

[taldcroft]

For a real-world example, reading the first Gaia ECSV source file here: https://cdn.gea.esac.esa.int/Gaia/gdr3/gaia_source/ (after unzipping) gives:

• fast io.ascii: 10.1 s (using the CSV reader just ignoring comments)
• pyarrow.csv: 1.5 s

[mhvk]

@taldcroft - this mostly looks good, and quite a few of my comments are really just nitpicks. Though I think I found a bug that more detailed tests would surely have found too...

p.s. The GAIA example shows this is really worth it, great!

[mhvk]

This comment is out of date and can be removed.

[taldcroft]

@mhvk - I added a suite of tests that covers all the read_csv arguments and a few more things.

A significant addition since you last looked is support for date, time, timestamp types. Currently these result in numpy datetime64 arrays, or an object array of datetime.time objects (for pure times like 12:34:45.233). An obvious idea is an option to allow converting to a Time mixin column where applicable.

Apart from that, I think this is mostly feature complete at least for this release. So it's just the docs, change log, what's new.

[mhvk]

This looks rather nice! One real comment left, with the suggestion to do the filling on the numpy side for speed and to save memory. But even that not serious.

A problem though is that most CI runs errored:

TypeError: ChunkedArray.to_numpy() takes no keyword arguments

EDIT: if this is a pyarrow version problem, one could probably set it to a higher minimum, since it is new for astropy.

[mhvk]

This is not needed, right?

[mhvk]

Are there any unsupported PyArrow types left? If not, the above logic is just for fill_value, which is not needed if there are no masked elements, so it could be moved inside the masked branch.

[taldcroft]

@mhvk - I think this ready for final review.

• Fixed the annoying Windows-only problems.
• Fixed a real bug in the comment handling and added a number of tests.
• Fixed most of the coverage issues. The remaining 3 bits of uncovered code are OK (import check and two data types where the code is obvious).
• Added a lot more documentation.

[taldcroft]

@mhvk - thanks so much for the review, as always this PR ended up in a much better place! I've addressed your minor docs comments and auto-merge squash is now enabled. 🤞 it gets in without another What's New merge conflict.