add pandas pyarrow backend support

[aaravind100]

Adds a subset of common data types from pyarrow data types.

fixes: #1262

[aaravind100]

There seems to be a bug in pandas.

this

[cosmicBboy]

thanks @aaravind100!

looks like there're failing tests:

FAILED tests/core/test_dtypes.py::test_check_equivalent[string[pyarrow]-'string[pyarrow]'] - assert False
 +  where False = <bound method DataType.check of DataType(string[pyarrow])>(DataType(string[pyarrow]))
 +    where <bound method DataType.check of DataType(string[pyarrow])> = DataType(string[pyarrow]).check
FAILED tests/core/test_pandas_engine.py::test_pandas_data_type[ArrowString] - AssertionError: assert DataType(string[pyarrow]) == DataType(string)
  
  Differing attributes:
  ['type']
  
  Drill down into differing attribute type:
    type: string[pyarrow] != string[pyarrow]
========== 2 failed, 1762 passed, 2 skipped, 1023 warnings in 52.38s ===========

You can use nox to try to reproduce the error, see this make target: https://github.com/unionai-oss/pandera/blob/main/Makefile#L55-L56

make nox-tests

You'll need mamba installed.

If you can run one specific test env:

nox -db mamba --envdir .nox-mamba -s "tests(extra='core', pydantic='1.10.11', python='3.11', pandas='2.2.0')"

[cosmicBboy]

There seems to be a bug in pandas.

This might be for backwards compatilibity reasons. I think when pandas < 2 supported pyarrow the string alias "string" would use pyarrow string dtype and "str" would use numpy objects dtype.

[aaravind100]

This might be for backwards compatilibity reasons. I think when pandas < 2 supported pyarrow the string alias "string" would use pyarrow string dtype and "str" would use numpy objects dtype.

Ah that makes sense.

Do you suggest fixing it with a condition?

[cosmicBboy]

yes, using a condition is okay.

also, we need to account for pandas < 2, see failures here: https://github.com/unionai-oss/pandera/actions/runs/9005141365/job/24742482507?pr=1628

I think the condition to define those classes need to check if a) pyarrow is installed and b) if pandas >= 2. You can use this function:

pandera/pandera/engines/utils.py

Line 12 in 63140c9

def pandas_version():

[aaravind100]

@cosmicBboy i ended up removing "string[pyarrow]" in the equivalents for ArrowString because pandas considers this astype("string[pyarrow]") as legacy string.

[aaravind100]

Hey @cosmicBboy, whats the best way to run the entire ci test suite in local? make nox-tests with the default config seem to chew through disk space 😅
I was using this pyspark image to do the development.

[cosmicBboy]

@aaravind100 you can simply do pytest tests to run the test suite for whatever your development environment is. Generally this will catch most errors.

For version specific tests, you can do nox --list to find a specific test combination to run and pass that into the -s flag, for example:

nox -db mamba --envdir .nox-mamba -s "tests(extra='core', pydantic='1.10.11', python='3.11', pandas='2.2.0')"

[aaravind100]

thanks, let me run it across a subset before pushing a commit.

This is what i meant by disk space usage 😄

du -h --max-depth=1 .
37K	./dev
191M	./tests
30K	./pandera.egg-info
6.0K	./.vscode
6.0K	./scripts
5.0M	./.git
18K	./__pycache__
54G	./.nox-mamba   # <--
1.7M	./.hypothesis
4.9M	./pandera
89K	./.pytest_cache
223K	./ci
18K	./asv_bench
502K	./docs
1.9M	./htmlcov
5.5M	./.mypy_cache
73K	./.github
54G	.

[cosmicBboy]

lol yeah I experience this too... yeah just nuke .nox-mamba and do the targeted runs

[cosmicBboy]

one strategy here is to do a check in the test_pandas_data_type function body and pytest.skip if pandas > 2 and pyarrow is installed.

[aaravind100]

i fixed it like this, i see this pattern in some other tests.

UNSUPPORTED_DTYPE_CLS: set[Any] = set()

# `string[pyarrow]` gets parsed to type `string` by pandas
if pandas_engine.PYARROW_INSTALLED and pandas_engine.PANDAS_2_0_0_PLUS:
    UNSUPPORTED_DTYPE_CLS.add(pandas_engine.ArrowString)


@pytest.mark.parametrize(
    "data_type",
    [
        data_type
        for data_type in pandas_engine.Engine.get_registered_dtypes()
        if data_type not in UNSUPPORTED_DTYPE_CLS
    ],
)
def test_pandas_data_type(data_type):
    ...

I just want to make sure it passes all the ci checks :D

[cosmicBboy]

lgtm

[aaravind100]

Ran the subset for python 3.10 and 3.11. Just this test for polars fails.

FAILED tests/polars/test_polars_container.py::test_dataframe_column_level_coerce - NotImplementedError

[cosmicBboy]

@aaravind100 if you rebase on the main branch this failing test has been fixed

[aaravind100]

@aaravind100 if you rebase on the main branch this failing test has been fixed

Thanks, it passes now. Pushing the commit now.

[cosmicBboy]

use typing.Set here for compat with py3.8

[aaravind100]

oops that was a force of habit, let me fix it

[codecov]

Codecov Report

Attention: Patch coverage is 99.02913% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 83.27%. Comparing base (4df61da) to head (2539773).
Report is 88 commits behind head on main.

Files	Patch %	Lines
pandera/engines/pandas_engine.py	99.02%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1628       +/-   ##
===========================================
- Coverage   94.29%   83.27%   -11.02%     
===========================================
  Files          91      116       +25     
  Lines        7024     8646     +1622     
===========================================
+ Hits         6623     7200      +577     
- Misses        401     1446     +1045     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[cosmicBboy]

this is amazing @aaravind100 ! all tests are passing (codecov is a false positive). gonna do one last review of the code over the weekend, but this is a huge feature set to support in pandera ❤️