Repository navigation
[Python] Build nested rows for struct values in pyarrow batch converters - #41201
Thanatos9404 wants to merge 3 commits into
Conversation
PyarrowBatchConverter.explode_batch built each element from pa.Scalar.as_py(), which returns struct values as dicts. With ReadFromParquet(as_rows=True), the element type declares those fields as nested rows, so code like row.address.city failed with AttributeError. PyarrowArrayBatchConverter had the same problem for row elements. Build the declared row type for struct values, including rows in optional fields and lists. A struct read with only some of its columns gets None for the others. Fixes apache#41100
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This PR updates PyarrowBatchConverter and PyarrowArrayBatchConverter to recursively reconstruct nested row types (RowTypeConstraint, NamedTuple, and @dataclass) from the dictionaries returned by PyArrow's Scalar.as_py(), including when nested inside optional or sequence fields. It also handles partial struct reads by defaulting omitted struct subfields to None and adds comprehensive unit and Parquet I/O tests.
🔍 General Feedback
- The recursive converter pre-computation in
_from_arrow_value_fnand_row_from_dict_fnis clean and avoids per-element type inspection overhead duringexplode_batch. - One edge case to address in
_from_arrow_value_fnis handling BeamTypeConstraintobjects (whichBatchConverter.from_typehintsproduces viatypehints.normalize()) and Python 3.10+ PEP 604 union types (X | None), wheretyping.get_origin()does not returntyping.Unionor standard collection origins.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
|
|
||
| origin = get_origin(typ) | ||
| args = get_args(typ) | ||
| if origin is Union: |
There was a problem hiding this comment.
🟠 High: _from_arrow_value_fn relies on typing.get_origin(typ) and origin is Union, which misses two cases:
- Normalized Beam
TypeConstraints:BatchConverter.from_typehintsnormalizeselement_typeviatypehints.normalize(element_type)before instantiatingPyarrowArrayBatchConverter(and Beam typehints likebeam.typehints.Optional[...]orbeam.typehints.List[...]can also appear in row fields). ForUnionConstraintandSequenceTypeConstraintsubclasses (ListConstraint,ABCSequenceTypeConstraint, etc.),typing.get_origin(typ)returnsNone, soBatchConverter.from_typehints(element_type=Optional[Address], batch_type=pa.Array)orBatchConverter.from_typehints(element_type=List[Address], batch_type=pa.Array)will haveself._convert = Noneand yield rawdicts instead ofAddressrows. - PEP 604 union syntax (
Address | None): On Python 3.10+,get_origin(Address | None)istypes.UnionTyperather thantyping.Union.
Consider converting non-row TypeConstraints via convert_to_python_type(typ) first (as schemas.py does in typing_to_runner_api) and checking origin in (Union, types.UnionType).
There was a problem hiding this comment.
The first case was real: with batch_type=pa.Array, Optional[Address] and List[Address] elements still came out as dicts because from_typehints normalizes the hint. a42a60f handles the normalized union and sequence constraints and adds a test for both. I matched the constraints directly since convert_to_python_type raises for Collection[...]. A top-level Address | None is normalized to the same union, so it is covered too; PEP 604 unions inside a row field fail earlier in schema conversion, which this PR does not touch.
|
Assigning reviewers: R: @shunping for label python. This pull request likely touches a core component ("core" label). Please review with scrutiny. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
…types BatchConverter.from_typehints normalizes the element type, so PyarrowArrayBatchConverter saw Beam UnionConstraint and sequence constraints, which typing.get_origin does not recognize, and returned struct values as dicts for Optional[Row] and List[Row] elements.
# Conflicts: # CHANGES.md
Fixes #41100
ReadFromParquet(as_rows=True)declares struct columns as nested rows in its element type, butPyarrowBatchConverter.explode_batchbuilt each element frompa.Scalar.as_py(), which returns structs as dicts. Sorow.address.city(or a Beam YAMLMapToFieldsexpression on a struct field) failed withAttributeError: 'dict' object has no attribute 'city'.explode_batchnow builds the declared row type for struct values, including rows in optional fields and in lists.PyarrowArrayBatchConverter.explode_batchhad the same problem for row elements and gets the same fix. Each converter works out once which fields need converting, so rows without nested structs take the same path as before.A struct read with only some of its columns (
columns=['address.city']) getsNonefor the fields that were not read.Behaviour change: nested struct values are rows now, so code that worked around this by indexing them as dicts (
row.address['city']) has to use attribute access. TheCHANGES.mdentry says so.Not changed here:
as_py()returns them.beam.Rowvalues (produce_batch) is [Bug]: [Python] WriteToParquet fails on schema'd PCollections whose nested struct fields hold beam.Row values #41101. Nested rows that Beam generates are NamedTuples, which pyarrow already accepts. I checked aReadFromParquet(as_rows=True)->Filter->WriteToParquetpipeline with a struct column and a list of structs, and it writes the same data back.Tests:
arrow_type_compatibility_test.py, for:pa.Arraybatch.parquetio_test.py. It reads a Parquet file with a struct column and a list of structs, and also does a partial struct read.arrow_type_compatibility_test.pyandbatch_test.pypass, along with the schema'd tests inparquetio_test.py.I also added a
CHANGES.mdentry under Bugfixes for 2.79.0.