feat: 2772 pyarrow doesnt permit selective reading with extensionarray

[tcawlfield]

Adding a module, awkward._connect.pyarrow_table_conv, to convert pyarrow tables with extensionarray types to ones without, storing metadata necessary to reconstruct the extensionarray types into the table-level metadata.

This allows reading select columns, as the tables in newly-written Parquet files will have native column types together with this metadata needed to convert back again on read.

[agoose77]

Super excited to see this work @tcawlfield!

[jpivarski]

It's looking good! The next step would be to make ak.to_arrow_table and ak.from_arrow_table (and, by extension, ak.to_parquet and ak.from_parquet) do this automatically. Then a lot of the existing tests in

• tests/test_1125_to_arrow_from_arrow.py
• tests/test_1154_arrow_tables_should_preserve_parameters.py
• tests/test_1294_to_and_from_parquet.py
• tests/test_1453_write_single_records_to_parquet.py
• tests/test_2340_unknown_type_to_arrow_and_back.py
• tests/test_2772_parquet_extn_array_metadata.py

would verify that the metadata is being fully preserved and arrays are being properly reconstructed.

I tried this:

>>> arrow_table = ak.to_arrow_table(   # tuples need metadata
...     ak.Array([[(1, 1.1), (2, 2.2)], [], [(3, 3.3)]]),
...     extensionarray=True,
... )
>>> ak._connect.pyarrow_table_conv.convert_awkward_arrow_table_to_native(
...     arrow_table
... )

because this function takes a pa.Table and returns a pa.Table, but it didn't work:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/jpivarski/irishep/awkward/src/awkward/_connect/pyarrow_table_conv.py", line 37, in convert_awkward_arrow_table_to_native
    return replace_schema(aatable, new_schema)
  File "/Users/jpivarski/irishep/awkward/src/awkward/_connect/pyarrow_table_conv.py", line 182, in replace_schema
    for col, new_field in zip(batch.itercolumns(), new_schema):
AttributeError: 'pyarrow.lib.RecordBatch' object has no attribute 'itercolumns'. Did you mean: 'num_columns'?

(I wanted to see what the collected metadata looks like.) Maybe I'm pulling the wrong function, though. (This whole new submodule consists of internal functions, so they don't have to have an obvious way to use them. They can be eclectic.)

This is looking good, but the real test will be when it's applied to all Awkward ↔ Arrow Table conversions, which will verify that the full conversions are lossless and correct.

[tcawlfield]

It's looking good! The next step would be to make ak.to_arrow_table and ak.from_arrow_table (and, by extension, ak.to_parquet and ak.from_parquet) do this automatically. Then a lot of the existing tests in

• tests/test_1125_to_arrow_from_arrow.py
• tests/test_1154_arrow_tables_should_preserve_parameters.py
• tests/test_1294_to_and_from_parquet.py
• tests/test_1453_write_single_records_to_parquet.py
• tests/test_2340_unknown_type_to_arrow_and_back.py
• tests/test_2772_parquet_extn_array_metadata.py

would verify that the metadata is being fully preserved and arrays are being properly reconstructed.

I tried this:

>>> arrow_table = ak.to_arrow_table(   # tuples need metadata
...     ak.Array([[(1, 1.1), (2, 2.2)], [], [(3, 3.3)]]),
...     extensionarray=True,
... )
>>> ak._connect.pyarrow_table_conv.convert_awkward_arrow_table_to_native(
...     arrow_table
... )

because this function takes a pa.Table and returns a pa.Table, but it didn't work:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/jpivarski/irishep/awkward/src/awkward/_connect/pyarrow_table_conv.py", line 37, in convert_awkward_arrow_table_to_native
    return replace_schema(aatable, new_schema)
  File "/Users/jpivarski/irishep/awkward/src/awkward/_connect/pyarrow_table_conv.py", line 182, in replace_schema
    for col, new_field in zip(batch.itercolumns(), new_schema):
AttributeError: 'pyarrow.lib.RecordBatch' object has no attribute 'itercolumns'. Did you mean: 'num_columns'?

(I wanted to see what the collected metadata looks like.) Maybe I'm pulling the wrong function, though. (This whole new submodule consists of internal functions, so they don't have to have an obvious way to use them. They can be eclectic.)

This is looking good, but the real test will be when it's applied to all Awkward ↔ Arrow Table conversions, which will verify that the full conversions are lossless and correct.

This should have worked. It works for me. Perhaps you found the first issue w/rt differing pyarrow versions? Or maybe I had failed to push local commits yesterday? I pushed some commits just now -- curious what this does for you now.

[tcawlfield]

I'm marking this as ready for review, even though there's failures on Python 3.8 still. I'll look into what it takes to reproduce this, but so far I'm hitting errors like:

AttributeError("'pyarrow.lib.LargeListType' object has no attribute 'field'")

that come specifically from here.

I'm having trouble installing all of requirements-test-minimal.txt in a Python 3.11 environment, although I got pyarrow 7.0.0 installed. And indeed LargeListArray in that version has a value_field property but not a field() method. I can probably special-case that with hasattr, but without reproducing the whole environment I'm not sure what would fail next. So I'll pause here and wait for any advice you might offer, and in the meantime try to get this working with an old Python somewhere.

You'll quickly notice that I finally broke down and converted ._connect.pyarrow into a package, with my new code as a submodule, .table_conv. I did a very minimal restructuring of the rest of pyarrow.py -- top-level __init__.py just handles safe importing with/without pyarrow and creates the same namespace of symbols regardless, so that other code could import anything here without having to get too fancy. The extension classes went into a separate submodule because that part was easy to chop off and it felt reasonable to do so. Everything else went into a grab-bag, .../pyarrow/conversions.py. (I'm happy to rename this.) I felt this was a safe place to draw the line between minor and major rearranging. My primary motivation was to pull the new table-conversion functions, convert_awkward_arrow_table_to_native and convert_native_arrow_table_to_awkward, up into ..._connect.pyarrow in a way that kept the new code separate, avoided circular imports, and worked well with/without pyarrow.

[jpivarski]

And indeed LargeListArray in that version has a value_field property but not a field() method. I can probably special-case that with hasattr, but without reproducing the whole environment I'm not sure what would fail next.

Actually, I remember that issue with LargeListArray. I might have even reported it, but it's probably from back when Arrow used JIRA instead of GitHub and I'd have trouble finding it. Special-casing with hasattr is definitely fine, and it's an isolated thing. There's no reason to expect a deluge of issues afterward.

Is the difficulty of making an environment with pyarrow 7 related to conda versus pip venv (because pyarrow has a binary dependency)? Making environments with the minimal and maximal dependencies works with conda (mamba).

Reorganizing all of the pyarrow connections into a directory is a good idea. Thanks!

[tcawlfield]

Okay I think this is all ready for review now, with tests passing!

Is the difficulty of making an environment with pyarrow 7 related to conda versus pip venv (because pyarrow has a binary dependency)? Making environments with the minimal and maximal dependencies works with conda (mamba).

My issue with pyarrow 7.0.0 is resolved. I'm not sure why I had trouble installing the older requirements earlier. It was a compilation error of some sort, but I was able to grab old pyarrow but latest versions of the other dependencies in our requirements-minimal.txt (something like that anyway) and work out all the changes needed. I left some earlier code commented-out with the rationale that we might in the future drop support for older pyarrow versions, and can clean up and/or improve the code. Some people don't like checking in commented-out code. "That's what git is for!" And this is true but I have things commented-out that might be worth reviewing in the future, or might also make intentions more clear. I do try to keep this to a bare minimum though.

[jpivarski]

This is looking great! I think this is ready to merge, if updating to main is uneventful (i.e. it merges cleanly and all tests continue to pass).