gemini-code-assist[bot] commented on code in PR #39236:
URL: https://github.com/apache/beam/pull/39236#discussion_r3623188482
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -2797,24 +2958,84 @@ def __exit__(self, *args):
pass
class ConvertToBeamRows(PTransform):
- def __init__(self, schema, dynamic_destinations, type_overrides=None):
+ def __init__(
+ self,
+ schema,
+ dynamic_destinations,
+ type_overrides=None,
+ schema_side_inputs=None):
self.schema = schema
self.dynamic_destinations = dynamic_destinations
self.type_overrides = type_overrides
+ self.schema_side_inputs = schema_side_inputs or ()
+
+ def _get_record_type_hint(self):
+ if callable(self.schema):
+ schema_hint = (
+ getattr(self.schema, '_union_schema', None) or
+ getattr(self.schema, '_table_schema', None) or
+ getattr(self.schema, '_beam_schema', None) or
+ getattr(self.schema, '_schema_hint', None) or
+ getattr(self.schema, '_output_types', None) or
+ getattr(self.schema, 'table_schema', None) or
+ getattr(self.schema, 'schema', None))
+ if schema_hint is not None:
+ if isinstance(
+ schema_hint,
+ (bigquery.TableSchema, bigquery.TableFieldSchema, str, dict)):
Review Comment:
Using `bigquery.TableSchema` and `bigquery.TableFieldSchema` inside
`bigquery.py` will raise a `NameError` at runtime. Please reference
`TableSchema` and `TableFieldSchema` directly.
```suggestion
(TableSchema, TableFieldSchema, str, dict)):
```
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -1956,6 +1993,120 @@ def _restore_table_ref(sharded_table_ref_elems_kv):
SCHEMA_AUTODETECT = 'SCHEMA_AUTODETECT'
+def dynamic_schema(schema_fn_or_map, union_schema=None):
+ """Helper to construct a dynamic schema callable with a union schema hint.
+
+ When using the BigQuery Storage Write API (`method=STORAGE_WRITE_API`) with
+ dynamic destinations, the cross-language transform requires a
PCollection-level
+ union schema containing all fields across all target tables for protobuf
+ serialization and type inference.
+
+ This helper provides the recommended best practice for constructing dynamic
+ schemas:
+
+ 1. **Dictionary Map**: If destination table schemas are provided as a
dictionary
+ mapping table names/specs to schemas (str, dict, or TableSchema), this
helper
+ automatically merges all fields into a single union schema.
+ 2. **Callable**: If a callable function is used, this helper attaches the
provided
+ `union_schema` to the callable as a schema hint (`_union_schema`).
+
+ Example using a dictionary map (union schema is auto-inferred)::
+
+ schema_map = {
+ 'project:dataset.users': 'id:INTEGER,name:STRING',
+ 'project:dataset.scores': 'id:INTEGER,score:INTEGER,active:BOOLEAN'
+ }
+
+ def get_destination(record):
+ if 'name' in record:
+ return 'project:dataset.users'
+ return 'project:dataset.scores'
+
+ schema_callable = dynamic_schema(schema_map)
+
+ elements | WriteToBigQuery(
+ table=get_destination,
+ method=WriteToBigQuery.Method.STORAGE_WRITE_API,
+ schema=schema_callable)
+
+ Example using a callable with explicit union schema::
+
+ def get_schema(destination, schema_side_input):
+ return schema_side_input[destination]
+
+ schema_callable = dynamic_schema(
+ get_schema,
+ union_schema='id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN')
+
+ elements | WriteToBigQuery(
+ table=get_destination,
+ method=WriteToBigQuery.Method.STORAGE_WRITE_API,
+ schema=schema_callable,
+ schema_side_inputs=(schema_side_input,))
+
+ Args:
+ schema_fn_or_map: A callable `(destination, *side_inputs) -> schema`
+ or a dictionary mapping destination strings to schemas (str, dict, or
+ TableSchema).
+ union_schema: (Optional) The union schema containing all fields across
+ target tables. Can be a string, dict, or TableSchema object. Required if
+ `schema_fn_or_map` is a callable.
+
+ Returns:
+ A callable with the attached `_union_schema` attribute for Storage Write
API.
+ """
+ if isinstance(schema_fn_or_map, dict):
+ if union_schema is None:
+ bq_schemas = [
+ bigquery_tools.get_bq_tableschema(s)
+ for s in schema_fn_or_map.values()
+ ]
+
+ def _merge_fields(field_a, field_b):
+ if field_a.type != field_b.type:
+ raise ValueError(
+ f"Conflicting types for field '{field_a.name}': "
+ f"{field_a.type} vs {field_b.type}")
+ if field_a.type in ('RECORD', 'STRUCT'):
+ merged_subfields = {}
+ for f in (field_a.fields or []):
+ merged_subfields[f.name] = f
+ for f in (field_b.fields or []):
+ if f.name in merged_subfields:
+ merged_subfields[f.name] = _merge_fields(
+ merged_subfields[f.name], f)
+ else:
+ merged_subfields[f.name] = f
+ field_a.fields = list(merged_subfields.values())
+ return field_a
+
+ merged_fields = {}
+ for schema in bq_schemas:
+ for field in schema.fields:
+ name = field.name
+ if name in merged_fields:
+ merged_fields[name] = _merge_fields(merged_fields[name], field)
+ else:
+ merged_fields[name] = field
+ union_schema = bigquery.TableSchema(fields=list(merged_fields.values()))
Review Comment:
Using `bigquery.TableSchema` inside `bigquery.py` will raise a `NameError`
at runtime because `bigquery` is the name of the module itself and is not
defined in the global namespace of this file. Please reference `TableSchema`
directly.
```suggestion
union_schema = TableSchema(fields=list(merged_fields.values()))
```
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -1956,6 +1993,120 @@ def _restore_table_ref(sharded_table_ref_elems_kv):
SCHEMA_AUTODETECT = 'SCHEMA_AUTODETECT'
+def dynamic_schema(schema_fn_or_map, union_schema=None):
+ """Helper to construct a dynamic schema callable with a union schema hint.
+
+ When using the BigQuery Storage Write API (`method=STORAGE_WRITE_API`) with
+ dynamic destinations, the cross-language transform requires a
PCollection-level
+ union schema containing all fields across all target tables for protobuf
+ serialization and type inference.
+
+ This helper provides the recommended best practice for constructing dynamic
+ schemas:
+
+ 1. **Dictionary Map**: If destination table schemas are provided as a
dictionary
+ mapping table names/specs to schemas (str, dict, or TableSchema), this
helper
+ automatically merges all fields into a single union schema.
+ 2. **Callable**: If a callable function is used, this helper attaches the
provided
+ `union_schema` to the callable as a schema hint (`_union_schema`).
+
+ Example using a dictionary map (union schema is auto-inferred)::
+
+ schema_map = {
+ 'project:dataset.users': 'id:INTEGER,name:STRING',
+ 'project:dataset.scores': 'id:INTEGER,score:INTEGER,active:BOOLEAN'
+ }
+
+ def get_destination(record):
+ if 'name' in record:
+ return 'project:dataset.users'
+ return 'project:dataset.scores'
+
+ schema_callable = dynamic_schema(schema_map)
+
+ elements | WriteToBigQuery(
+ table=get_destination,
+ method=WriteToBigQuery.Method.STORAGE_WRITE_API,
+ schema=schema_callable)
+
+ Example using a callable with explicit union schema::
+
+ def get_schema(destination, schema_side_input):
+ return schema_side_input[destination]
+
+ schema_callable = dynamic_schema(
+ get_schema,
+ union_schema='id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN')
+
+ elements | WriteToBigQuery(
+ table=get_destination,
+ method=WriteToBigQuery.Method.STORAGE_WRITE_API,
+ schema=schema_callable,
+ schema_side_inputs=(schema_side_input,))
+
+ Args:
+ schema_fn_or_map: A callable `(destination, *side_inputs) -> schema`
+ or a dictionary mapping destination strings to schemas (str, dict, or
+ TableSchema).
+ union_schema: (Optional) The union schema containing all fields across
+ target tables. Can be a string, dict, or TableSchema object. Required if
+ `schema_fn_or_map` is a callable.
+
+ Returns:
+ A callable with the attached `_union_schema` attribute for Storage Write
API.
+ """
+ if isinstance(schema_fn_or_map, dict):
+ if union_schema is None:
+ bq_schemas = [
+ bigquery_tools.get_bq_tableschema(s)
+ for s in schema_fn_or_map.values()
+ ]
Review Comment:
The `_merge_fields` helper mutates the fields of the input schemas in-place.
If the user passes existing `TableSchema` objects in `schema_fn_or_map`,
`get_bq_tableschema` returns them directly, meaning this helper will silently
mutate the user's original schema objects. To prevent unexpected side effects,
please deepcopy the schemas before merging them.
```suggestion
import copy
bq_schemas = [
copy.deepcopy(bigquery_tools.get_bq_tableschema(s))
for s in schema_fn_or_map.values()
]
```
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -2797,24 +2958,84 @@ def __exit__(self, *args):
pass
class ConvertToBeamRows(PTransform):
- def __init__(self, schema, dynamic_destinations, type_overrides=None):
+ def __init__(
+ self,
+ schema,
+ dynamic_destinations,
+ type_overrides=None,
+ schema_side_inputs=None):
self.schema = schema
self.dynamic_destinations = dynamic_destinations
self.type_overrides = type_overrides
+ self.schema_side_inputs = schema_side_inputs or ()
+
+ def _get_record_type_hint(self):
+ if callable(self.schema):
+ schema_hint = (
+ getattr(self.schema, '_union_schema', None) or
+ getattr(self.schema, '_table_schema', None) or
+ getattr(self.schema, '_beam_schema', None) or
+ getattr(self.schema, '_schema_hint', None) or
+ getattr(self.schema, '_output_types', None) or
+ getattr(self.schema, 'table_schema', None) or
+ getattr(self.schema, 'schema', None))
+ if schema_hint is not None:
+ if isinstance(
+ schema_hint,
+ (bigquery.TableSchema, bigquery.TableFieldSchema, str, dict)):
+ row_type_hints =
bigquery_tools.get_beam_typehints_from_tableschema(
+ schema_hint, self.type_overrides)
+ return RowTypeConstraint.from_fields(row_type_hints)
+ elif isinstance(schema_hint, RowTypeConstraint):
+ return schema_hint
+ return RowTypeConstraint.from_fields([])
+ else:
+ row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema(
+ self.schema, self.type_overrides)
+ return RowTypeConstraint.from_fields(row_type_hints)
def expand(self, input_dicts):
if self.dynamic_destinations:
- return (
- input_dicts
- | "Convert dict to Beam Row" >> beam.Map(
- lambda row, schema=DoFn.SetupContextParam(
- StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args=
- [self.schema]): beam.Row(
- **{
- StorageWriteToBigQuery.DESTINATION: row[0],
- StorageWriteToBigQuery.RECORD: bigquery_tools.
- beam_row_from_dict(row[1], schema)
- })))
+ if callable(self.schema):
+ record_hint = self._get_record_type_hint()
+ union_field_names = [
+ name for name, _ in getattr(record_hint, '_fields', ())
+ ]
+
+ def convert_dynamic_row(row, *schema_side_inputs):
+ dest, dict_row = row[0], row[1]
+ record_schema = self.schema(dest, *schema_side_inputs)
+ record_row = bigquery_tools.beam_row_from_dict(
+ dict_row, record_schema)
+ if union_field_names:
+ record_dict = record_row._asdict()
+ record_row = beam.Row(
+ **{
+ name: record_dict.get(name, None)
+ for name in union_field_names
+ })
+ return beam.Row(
+ **{
+ StorageWriteToBigQuery.DESTINATION: dest,
+ StorageWriteToBigQuery.RECORD: record_row
+ })
Review Comment:
Constructing a new `beam.Row` using `beam.Row(**...)` with dynamic keyword
arguments causes Beam to infer the schema on the fly. For fields that are
`None`, their types will be inferred as `NoneType` (or `Any`), which means
different rows will have different inferred schemas. This will cause
serialization/encoding errors (e.g., `CoderException`) at the cross-language
boundary because the `RowCoder` expects a consistent schema matching the union
schema.
Additionally, converting to a dict, performing a dict comprehension, and
instantiating a new `beam.Row` for every single element introduces significant
CPU overhead.
Instead, you can use `bigquery_tools.beam_row_from_dict` with the
pre-resolved `union_bq_schema` to safely and efficiently pad/convert the row to
the union schema.
```python
union_schema = getattr(self.schema, '_union_schema', None)
union_bq_schema = (
bigquery_tools.get_bq_tableschema(union_schema)
if union_schema else None)
def convert_dynamic_row(row, *schema_side_inputs):
dest, dict_row = row[0], row[1]
record_schema = self.schema(dest, *schema_side_inputs)
record_row = bigquery_tools.beam_row_from_dict(
dict_row, record_schema)
if union_bq_schema:
record_row = bigquery_tools.beam_row_from_dict(
record_row._asdict(), union_bq_schema)
return beam.Row(
**{
StorageWriteToBigQuery.DESTINATION: dest,
StorageWriteToBigQuery.RECORD: record_row
})
```
--
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]
