vishalmore90 opened a new issue, #41238:
URL: https://github.com/apache/beam/issues/41238
### What happened?
In Apache Beam's YAML SDK, `YamlProviders.Flatten._merge_schemas`
(`sdks/python/apache_beam/yaml/yaml_provider.py:1026-1066`) unifies schemas
across input `PCollection`s. However, the implementation has three critical
flaws:
#### 1. Non-deterministic Field Ordering Across Runs and Workers
In `sdks/python/apache_beam/yaml/yaml_provider.py:1050-1064`:
```python
# Merge all field names and types.
all_field_names = set().union(*(s.keys() for s in schemas))
unified_fields = {}
for name in all_field_names:
...
unified_fields[name] = ...
# Create unified schema
if unified_fields:
unified_schema = named_fields_to_schema(list(unified_fields.items()))
return named_tuple_from_schema(unified_schema)
```
Because `all_field_names` is an unordered Python `set`, iterating over it
(`for name in all_field_names`) depends on hash randomization
(`PYTHONHASHSEED`). Across different runs, workers, or Python environments, the
fields in `unified_schema` and the generated NamedTuple class are ordered
non-deterministically.
For example, when input schemas have fields `['first', 'second', 'third',
'fourth', 'fifth']`:
- With `PYTHONHASHSEED=1`: fields are ordered `['fifth', 'first', 'third',
'fourth', 'second']`
- With `PYTHONHASHSEED=42`: fields are ordered `['third', 'fifth', 'fourth',
'second', 'first']`
- With `PYTHONHASHSEED=100`: fields are ordered `['first', 'fifth',
'second', 'fourth', 'third']`
In Beam schemas and Python `namedtuple`s, field ordering determines
positional indexing (`row[0]`, `row[1]`), schema compatibility checks, and
column ordering in schema-aware sinks (e.g., BigQuery, Parquet, Avro).
Non-deterministic ordering causes schema mismatch errors and inconsistent sink
schemas.
#### 2. Loss of Concrete Type: `Optional[T]` Merged with `T` Degrades to
`Optional[Any]`
In `sdks/python/apache_beam/yaml/yaml_provider.py:1052-1057`:
```python
for name in all_field_names:
present_types = {s[name] for s in schemas if name in s}
if len(present_types) > 1:
unified_fields[name] = Optional[Any]
else:
unified_fields[name] = Optional[present_types.pop()]
```
In Python typing, `int != Optional[int]` (since `Optional[int]` is
`Union[int, NoneType]`). If one input `PCollection` defines a field as `int`
and another input defines it as `Optional[int]`, `present_types` becomes `{int,
Optional[int]}`, which has length 2.
Consequently, `len(present_types) > 1` triggers, and the field's type is
degraded to `Optional[Any]`, completely discarding the concrete `int` type
constraint even though the type is identical across all inputs.
#### 3. Unconditional Wrapping in `Optional[...]` and Performance Overhead
on Identical Schemas
Even when all input `PCollection`s share the exact same schema with
non-optional, required fields:
- Every field is unconditionally wrapped in `Optional[...]` (e.g. `int`
becomes `Optional[int]`).
- `_merge_schemas` returns a newly synthesized schema rather than
recognizing that schemas are already uniform.
- `Flatten.expand` applies `beam.Map(_unify_element_with_schema)` to every
input collection, forcing every single row in the pipeline to be converted from
NamedTuple to `dict` and back to NamedTuple, incurring substantial per-element
CPU and memory overhead instead of directly delegating to native `beam.Flatten`.
---
### Steps to Reproduce
Run the following minimal script:
```python
from typing import Optional, Any
from apache_beam.yaml.yaml_provider import YamlProviders
# Case 1: Combining int and Optional[int]
s1 = {"id": int, "name": str}
s2 = {"id": Optional[int], "name": str}
all_field_names = set().union(*(s.keys() for s in [s1, s2]))
unified_fields = {}
for name in all_field_names:
present_types = {s[name] for s in [s1, s2] if name in s}
if len(present_types) > 1:
unified_fields[name] = Optional[Any]
else:
unified_fields[name] = Optional[present_types.pop()]
print("Unified 'id' type:", unified_fields["id"])
# Output: typing.Optional[typing.Any] (Expected: typing.Optional[int])
```
---
### Proposed Solution
1. **Preserve Deterministic Field Ordering**:
Instead of `set().union(*(s.keys() for s in schemas))`, preserve
insertion order across input schemas using `dict.fromkeys`:
```python
all_field_names = list(dict.fromkeys(name for s in schemas for name in
s.keys()))
```
2. **Unwrap `Optional` Before Comparing Types**:
Extract inner types using
`apache_beam.typehints.trivial_inference.extract_optional_type` (or
`typing.get_args`) so that `{int, Optional[int]}` resolves to `int` and unifies
cleanly as `Optional[int]`.
3. **Only Mark Fields Optional When Necessary**:
A field only needs to be `Optional` if it was already optional in an
input schema or is missing from at least one input schema (`len([s for s in
schemas if name in s]) < len(schemas)`).
4. **Bypass Unification When Schemas Are Already Identical**:
If all input schemas are identical, return `None` from `_merge_schemas`
so `Flatten.expand` uses native `beam.Flatten(**pipeline_arg)` with zero
overhead.
---
### Issue Priority
Priority: 2 (default / most bugs should be filed as P2)
### Issue Components
- [x] Component: Beam YAML
- [x] Component: Python SDK
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
