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?

[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_.

[rok]

Another rebase on master.

[rok]

@pitrou could you do a quick pass here in case anything stands out please?

[pitrou]

Side note: it would be nice if one could write json_array = pa.array(arr, json_type).
Perhaps open a feature request?

[rok]

That would be a nice feature. #44406

[pitrou]

Pity this doesn't raise TypeError.

[rok]

We could catch and raise it but it's probably not a good idea.

[pitrou]

Indeed. TypeError would have to be raised at the C++ level instead. Anyway, this is out of scope for this PR.

[rok]

@pitrou do you think this merits a merge or should we wait for another review?

[pitrou]

What's the plan for the Parquet arrow_extensions_enabled option?

[rok]

@pitrou

What's the plan for the Parquet arrow_extensions_enabled option?

Perhaps we should open another issue for it? Current implementation seems to roundtrip to parquet ok.
I'd propose something like this:

diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd
index d6aebd8284..32e2618ecf 100644
--- a/python/pyarrow/_parquet.pxd
+++ b/python/pyarrow/_parquet.pxd
@@ -405,6 +405,7 @@ cdef extern from "parquet/api/reader.h" namespace "parquet" nogil:
         CCacheOptions cache_options() const
         void set_coerce_int96_timestamp_unit(TimeUnit unit)
         TimeUnit coerce_int96_timestamp_unit() const
+        void set_arrow_extensions_enabled(c_bool enabled)
 
     ArrowReaderProperties default_arrow_reader_properties()
 
diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx
index 254bfe3b09..6ae1726c71 100644
--- a/python/pyarrow/_parquet.pyx
+++ b/python/pyarrow/_parquet.pyx
@@ -1441,7 +1441,8 @@ cdef class ParquetReader(_Weakrefable):
              FileDecryptionProperties decryption_properties=None,
              thrift_string_size_limit=None,
              thrift_container_size_limit=None,
-             page_checksum_verification=False):
+             page_checksum_verification=False,
+             arrow_extensions_enabled=False):
         """
         Open a parquet file for reading.
 
@@ -1458,6 +1459,7 @@ cdef class ParquetReader(_Weakrefable):
         thrift_string_size_limit : int, optional
         thrift_container_size_limit : int, optional
         page_checksum_verification : bool, default False
+        arrow_extensions_enabled: bool, default False
         """
         cdef:
             shared_ptr[CFileMetaData] c_metadata
@@ -1522,6 +1524,9 @@ cdef class ParquetReader(_Weakrefable):
         if read_dictionary is not None:
             self._set_read_dictionary(read_dictionary, &arrow_props)
 
+        if arrow_extensions_enabled:
+            arrow_props.set_arrow_extensions_enabled(<c_bool>True)
+
         with nogil:
             check_status(builder.memory_pool(self.pool)
                          .properties(arrow_props)

[pitrou]

@rok, yes, we should open a new issue for it

[rok]

Opened an issue for the arrow_extensions_enabled #44500 problem.

[conbench-apache-arrow]

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

There were no benchmark performance regressions. 🎉

The full Conbench report has more details.