fix: avoid "Unable to determine type" warning with JSON columns in to_dataframe

[tswast]

Based on #2144, which should merge first.

TODO:

• tests
• maybe we don't want string columns for JSON? Update: keeping string because that's consistent with (1) current behavior and (2) the BQ Storage Read API.

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

• Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
• Ensure the tests and linter pass
• Code coverage does not decrease (if any source code was changed)
• Appropriate docs were updated (if necessary)

Fixes #1580
🦕

[tswast]

I might actually want to do something in db-dtypes so that even though it's a string the unboxed version would give a parsed object like the behavior is when the REST API is used.

[tswast]

Right now the behavior is inconsistent across REST and BQ Storage API.

[tswast]

Marking as do not merge for now, as this makes JSON dtype consistent now but always return string dtype like the BQ Storage Read API code path does, which isn't ideal.

[tswast]

Actually, I think this needs a few more tests. I'm testing manually with pytest 'tests/system/test_to_gbq.py::test_dataframe_round_trip_with_table_schema[load_csv-json]' from googleapis/python-bigquery-pandas#893, but it's currently failing because we parse the JSON string in _row_iterator_page_columns, but we actually want to keep those as strings to use the json_ pyarrow type.

[tswast]

This is the key change. Mostly aligns with bigframes, but we've left off pyarrow.json_(pyarrow.string()) because of pandas-dev/pandas#60958.

[tswast]

Marking as do not merge again. I'll split out the refactor into a separate PR first.

Edit: Mailed #2144

[tswast]

I've added regression tests for #1580

[tswast]

Test failure:

____________________ test_to_arrow_query_with_empty_results ____________________

bigquery_client = 

    def test_to_arrow_query_with_empty_results(bigquery_client):
        """
        JSON regression test for https://github.com/googleapis/python-bigquery/issues/1580.
        """
        job = bigquery_client.query(
            """
            select
            123 as int_col,
            '' as string_col,
            to_json('{}') as json_col,
            struct(to_json('[]') as json_field, -1 as int_field) as struct_col,
            [to_json('null')] as json_array_col,
            from unnest([])
            """
        )
        table = job.to_arrow()
        assert list(table.column_names) == [
            "int_col",
            "string_col",
            "json_col",
            "struct_col",
            "json_array_col",
        ]
        assert table.shape == (0, 5)
>       assert list(table.field("struct_col").type.names) == ["json_field", "int_field"]
E       AttributeError: 'pyarrow.lib.StructType' object has no attribute 'names'

Need to update this to support older pyarrow.

[GarrettWu]

Mapping to pa.string won't achieve round-trip? Meaning a value saving to local won't be able to be identified as JSON back. Does it matter to bigframes?

[tswast]

BigQuery sets metadata on the Field that can be used to determine this type. I don't want to diverge from BigQuery Storage Read API behavior.

In bigframes and pandas-gbq, we have the BigQuery schema available to disambiguate to customize the pandas types.

[tswast]

I can investigate if such an override is also possible here.

[tswast]

I made some progress plumbing through a json_type everywhere it would need to go to be able to override this, but once I got to to_arrow_iterable, it kinda breaks down. There we very much just return the pages we get from BQ Storage Read API. I don't really want to override that, as it adds new layers of complexity to what was a relatively straightforward internal API.

I'd prefer to leave this as-is without the change to allow overriding the arrow type.

[tswast]

If there's still objections, I can try the same but just with the pandas data type. That gets a bit awkward when it comes to struct, though.

[GarrettWu]

Thanks for digging into it.

[tswast]

It might be important to have this json_arrow_type feature for the work Chelsea is doing in bigframes and pandas-gbq. I'll give this another try, but I think since it's a feature it should be a separate PR.

[tswast]

Started #2149 but I'm hitting some roadblocks. Will pause for now.

[chalmerlowe]

#QUESTION:

The test suite for test_to_arrow_query_with_empty_results is more robust than this one.
Is there a reason for the difference?

[tswast]

We're using object dtype in pandas for STRUCT and ARRAY columns right now. This means there's not much we can inspect about subfields/subtypes like we can with Arrow.

Related: I'd like to introduce a dtype_backend parameter like pandas has to pandas-gbq: googleapis/python-bigquery-pandas#621

In the meantime, I tend to use:

df = (
    results
    .to_arrow(bqstorage_client=bqstorage_client)
    .to_pandas(types_mapper=lambda type_ : pandas.ArrowDtype(type_))
)

instead of to_dataframe() in my personal projects, as this will give me the more structured Arrow types all the time.

[chalmerlowe]

LGTM.

Included a question for my own edification.
Not a blocker.

Approved.