GH-44066: [Python] Add Python wrapper for JsonExtensionType

[rok]

Rationale for this change

We added canonical JsonExtensionType and we should make it usable from Python.

What changes are included in this PR?

Python wrapper for JsonExtensionType and JsonArray are added on Python side as well as JsonArray on c++ side.

Are these changes tested?

Python tests for the extension type and array are included.

Are there any user-facing changes?

This adds a json canonical extension type to pyarrow.

• GitHub Issue: [Python] Add Python wrapper for JsonExtensionType #44066

[github-actions]

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

[jorisvandenbossche]

Thanks for working on this!

Added some inline comments. One other comment:

• Can you add the new classes to test_extension_type_constructor_errors in test_misc.py

Are you planning to work on enabling the parquet integration in a later PR?

[jorisvandenbossche]

Should we maybe note here that this array class does not guarantee that the data is actually valid JSON? (which I assume is the case?)

For example in the code example below, you can in theory put whatever string data you want in the array.

[rok]

Added a note. Let me know if it's not extensive enough.

[jorisvandenbossche]

Can you also check equality at the type instance level?

Suggested change

	assert json_type.__class__ is pa.JsonType
	assert json_type.__class__ is pa.JsonType
	assert json_type == pa.json(storage_type)
	assert json_type != storage_type

Which also makes me wonder: if the storage type is different (eg string vs large string), does it then return False? (maybe also good to test that)

[rok]

Thanks for working on this!

Thanks for the review!

• Can you add the new classes to test_extension_type_constructor_errors in test_misc.py

Added.

Are you planning to work on enabling the parquet integration in a later PR?

Do we need to do something Python-side? C++ side was covered by #13901. I've added a basic parquet test here and it works for me locally.

Question: any idea why pytest is saying pa.json is not callable in CI? Locally this works as expected.

[jorisvandenbossche]

Do we need to do something Python-side? C++ side was covered by #13901. I've added a basic parquet test here and it works for me locally.

I would be surprised if it works locally out of the box because IIRC the option was disabled by default in C++ for now? (so what we would need on the python side is expose that new arrow_extensions_enabled options so that the user can enable it)
Although I suppose that you wrote the file with pyarrow, and at that point we give priority to the stored schema, I suppose, even if the option is not enabled.

Question: any idea why pytest is saying pa.json is not callable in CI? Locally this works as expected.

There might somehow we some conflict with the stdlib module?

[jorisvandenbossche]

Question: any idea why pytest is saying pa.json is not callable in CI? Locally this works as expected.

There might somehow we some conflict with the stdlib module?

Actually not that, but we have a pyarrow.json module already. So we will have to call this function differently anyway. For others with a conflict with added a trailing underscore, so could do that here as well.

[jorisvandenbossche]

Maybe move those tests to parquet subdir? (e.g. python/pyarrow/tests/parquet/test_data_types.py)
I know we already have parquet related tests in this file, but those are for custom extension type support, while this will be a built-in extension type

[rok]

Moved.

[rok]

Do we need to do something Python-side? C++ side was covered by #13901. I've added a basic parquet test here and it works for me locally.

I would be surprised if it works locally out of the box because IIRC the option was disabled by default in C++ for now? (so what we would need on the python side is expose that new arrow_extensions_enabled options so that the user can enable it) Although I suppose that you wrote the file with pyarrow, and at that point we give priority to the stored schema, I suppose, even if the option is not enabled.

I see. I assumed arrow_extensions_enabled was on by default. Do we want to expose a global setter? E.g.: pyarrow.parquet.set_arrow_extensions_enabled? Or would it be preferred to add a parameter for pq.write_table?

Question: any idea why pytest is saying pa.json is not callable in CI? Locally this works as expected.

There might somehow we some conflict with the stdlib module?
[..]
Actually not that, but we have a pyarrow.json module already. So we will have to call this function differently anyway. For > others with a conflict with added a trailing underscore, so could do that here as well.

It seems this was the case, changed extension type name to pa.json_.

[felipecrv]

Nice!

[rok]

Thanks for providing that :)

[rok]

@jorisvandenbossche ping :)