sdks/python/apache_beam/yaml/yaml_io.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""This module contains the Python implementations for the builtin IOs.

They are referenced from standard_io.py.

Note that in the case that they overlap with other (likely Java)
implementations of the same transforms, the configs must be kept in sync.
"""

import io
import json
from collections.abc import Callable
from collections.abc import Iterable
from collections.abc import Mapping
from typing import Any
from typing import Optional
from typing import Union

import fastavro

import apache_beam as beam
import apache_beam.io as beam_io
from apache_beam import coders
from apache_beam.coders.row_coder import RowCoder
from apache_beam.io import ReadFromBigQuery
from apache_beam.io import ReadFromTFRecord
from apache_beam.io import WriteToBigQuery
from apache_beam.io import WriteToTFRecord
from apache_beam.io import avroio
from apache_beam.io import mongodbio
from apache_beam.io.filesystem import CompressionTypes
from apache_beam.io.gcp.bigquery import BigQueryDisposition
from apache_beam.portability.api import schema_pb2
from apache_beam.typehints import schemas
from apache_beam.typehints.row_type import RowTypeConstraint
from apache_beam.utils.timestamp import Timestamp
from apache_beam.yaml import json_utils
from apache_beam.yaml import yaml_errors
from apache_beam.yaml import yaml_provider
from apache_beam.yaml import yaml_utils


def read_from_text(path: str):
  # TODO(yaml): Consider passing the filename and offset, possibly even
  # by default.

  """Reads lines from a text files.

  The resulting PCollection consists of rows with a single string field named
  "line."

  Args:
    path (str): The file path to read from.  The path can contain glob
      characters such as ``*`` and ``?``.
  """
  return beam_io.ReadFromText(path) | beam.Map(lambda s: beam.Row(line=s))


@beam.ptransform_fn
def write_to_text(pcoll, path: str):
  """Writes a PCollection to a (set of) text files(s).

  The input must be a PCollection whose schema has exactly one field.

  Args:
      path (str): The file path to write to. The files written will
        begin with this prefix, followed by a shard identifier.
  """
  try:
    field_names = [
        name for name, _ in schemas.named_fields_from_element_type(
            pcoll.element_type)
    ]
  except Exception as exn:
    raise ValueError(
        "WriteToText requires an input schema with exactly one field.") from exn
  if len(field_names) != 1:
    raise ValueError(
        "WriteToText requires an input schema with exactly one field, got %s" %
        field_names)
  sole_field_name, = field_names
  return pcoll | beam.Map(
      lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path)


def read_from_bigquery(
    *,
    table: Optional[str] = None,
    query: Optional[str] = None,
    row_restriction: Optional[str] = None,
    fields: Optional[Iterable[str]] = None,
    schema: Optional[Any] = None):
  """Reads data from BigQuery.

  Exactly one of table or query must be set.
  If query is set, neither row_restriction nor fields should be set.

  Args:
    table (str): The table to read from, specified as `DATASET.TABLE`
      or `PROJECT:DATASET.TABLE`.
    query (str): A query to be used instead of the table argument.
    row_restriction (str): Optional SQL text filtering statement, similar to a
      WHERE clause in a query. Aggregates are not supported. Restricted to a
      maximum length for 1 MB.
    selected_fields (list[str]): Optional List of names of the fields in the
      table that should be read. If empty, all fields will be read. If the
      specified field is a nested field, all the sub-fields in the field will be
      selected. The output field order is unrelated to the order of fields
      given here.
    schema (dict): Required when query is set. A BigQuery schema describing
      the query result columns, e.g.
      ``{'fields': [{'name': 'col', 'type': 'STRING', 'mode': 'NULLABLE'}]}``.
      Not applicable when reading from a table (schema is auto-derived).
  """
  if query is None:
    assert table is not None
  else:
    assert table is None and row_restriction is None and fields is None
    if schema is None:
      raise ValueError(
          "When using 'query' in ReadFromBigQuery YAML transform, "
          "'schema' is required to define the output row structure.")
  return ReadFromBigQuery(
      query=query,
      table=table,
      row_restriction=row_restriction,
      selected_fields=fields,
      method='DIRECT_READ',
      output_type='BEAM_ROW',
      query_output_schema=schema)


def write_to_bigquery(
    table: str,
    *,
    create_disposition: Optional[str] = BigQueryDisposition.CREATE_IF_NEEDED,
    write_disposition: Optional[str] = BigQueryDisposition.WRITE_APPEND,
    error_handling=None):
  f"""Writes data to a BigQuery table.

  Args:
    table (str): The table to read from, specified as `DATASET.TABLE`
      or `PROJECT:DATASET.TABLE`.
      create_disposition (BigQueryDisposition): A string describing what
        happens if the table does not exist. Possible values are:

        * :attr:`{BigQueryDisposition.CREATE_IF_NEEDED}`: create if does not
          exist.
        * :attr:`{BigQueryDisposition.CREATE_NEVER}`: fail the write if does not
          exist.

        Defaults to `{BigQueryDisposition.CREATE_IF_NEEDED}`.

      write_disposition (BigQueryDisposition): A string describing what happens
        if the table has already some data. Possible values are:

        * :attr:`{BigQueryDisposition.WRITE_TRUNCATE}`: delete existing rows.
        * :attr:`{BigQueryDisposition.WRITE_APPEND}`: add to existing rows.
        * :attr:`{BigQueryDisposition.WRITE_EMPTY}`: fail the write if table not
          empty.

        For streaming pipelines WriteTruncate can not be used.

        Defaults to `{BigQueryDisposition.WRITE_APPEND}`.

      error_handling: If specified, should be a mapping giving an output into
        which to emit records that failed to be written to BigQuery, as
        described at https://beam.apache.org/documentation/sdks/yaml-errors/
        Otherwise permanently failing records will cause pipeline failure.
        Note: error_handling requires the Storage Write API method and is not
        supported with File Loads (FILE_LOADS method). When error_handling is
        specified, the transform will automatically use STORAGE_WRITE_API.
  """
  class WriteToBigQueryHandlingErrors(beam.PTransform):
    def default_label(self):
      return 'WriteToBigQuery'

    def expand(self, pcoll):
      write_result = pcoll | WriteToBigQuery(
          table,
          method=WriteToBigQuery.Method.STORAGE_WRITE_API
          if error_handling else None,
          create_disposition=create_disposition,
          write_disposition=write_disposition,
          temp_file_format='AVRO')
      if error_handling and 'output' in error_handling:
        # TODO: Support error rates.
        return {
            'post_write': write_result.failed_rows_with_errors
            | beam.FlatMap(lambda x: None),
            error_handling['output']: write_result.failed_rows_with_errors
        }
      else:
        if write_result._method == WriteToBigQuery.Method.FILE_LOADS:
          # Never returns errors, just fails.
          return {
              'post_write': write_result.destination_load_jobid_pairs
              | beam.FlatMap(lambda x: None)
          }
        else:

          # This should likely be pushed into the BQ read itself to avoid
          # the possibility of silently ignoring errors.
          def raise_exception(failed_row_with_error):
            raise RuntimeError(failed_row_with_error.error_message)

          _ = write_result.failed_rows_with_errors | beam.Map(raise_exception)
          return {
              'post_write': write_result.failed_rows_with_errors
              | beam.FlatMap(lambda x: None)
          }

  return WriteToBigQueryHandlingErrors()


def _create_parser(
    format,
    schema: Any) -> tuple[schema_pb2.Schema, Callable[[bytes], beam.Row]]:

  format = format.upper()

  def _validate_schema():
    if not schema:
      raise ValueError(
          f'{format} format requires valid {format} schema to be passed to '
          f'schema parameter.')

  if format == 'RAW':
    if schema:
      raise ValueError('RAW format does not take a schema')
    return (
        schema_pb2.Schema(fields=[schemas.schema_field('payload', bytes)]),
        lambda payload: beam.Row(payload=payload))
  if format == 'STRING':
    if schema:
      raise ValueError('STRING format does not take a schema')
    return (
        schema_pb2.Schema(fields=[schemas.schema_field('payload', str)]),
        lambda payload: beam.Row(payload=payload.decode('utf-8')))
  elif format == 'JSON':
    _validate_schema()
    beam_schema = json_utils.json_schema_to_beam_schema(schema)
    return beam_schema, json_utils.json_parser(beam_schema, schema)
  elif format == 'AVRO':
    _validate_schema()
    beam_schema = avroio.avro_schema_to_beam_schema(schema)
    covert_to_row = avroio.avro_dict_to_beam_row(schema, beam_schema)
    return (
        beam_schema, lambda record: covert_to_row(
            fastavro.schemaless_reader(io.BytesIO(record), schema)))
  elif format == 'PROTO':
    _validate_schema()
    beam_schema = json_utils.json_schema_to_beam_schema(schema)
    return beam_schema, RowCoder(beam_schema).decode
  else:
    raise ValueError(f'Unknown format: {format}')


def _create_formatter(
    format, schema: Any,
    beam_schema: schema_pb2.Schema) -> Callable[[beam.Row], bytes]:

  if format.islower():
    format = format.upper()

  if format == 'RAW':
    if schema:
      raise ValueError('RAW format does not take a schema')
    field_names = [field.name for field in beam_schema.fields]
    if len(field_names) != 1:
      raise ValueError(f'Expecting exactly one field, found {field_names}')

    def convert_to_bytes(row):
      output = getattr(row, field_names[0])
      if isinstance(output, bytes):
        return output
      elif isinstance(output, str):
        return output.encode('utf-8')
      else:
        raise ValueError(
            f"Cannot encode payload for WriteToPubSub. "
            f"Expected valid string or bytes object, "
            f"got {repr(output)} of type {type(output)}.")

    return convert_to_bytes
  elif format == 'JSON':
    return json_utils.json_formater(beam_schema)
  elif format == 'AVRO':
    avro_schema = schema or avroio.beam_schema_to_avro_schema(beam_schema)
    from_row = avroio.beam_row_to_avro_dict(avro_schema, beam_schema)

    def formatter(row):
      buffer = io.BytesIO()
      fastavro.schemaless_writer(buffer, avro_schema, from_row(row))
      buffer.seek(0)
      return buffer.read()

    return formatter
  elif format == 'PROTO':
    return RowCoder(beam_schema).encode
  else:
    raise ValueError(f'Unknown format: {format}')


@beam.ptransform_fn
@yaml_errors.maybe_with_exception_handling_transform_fn
def read_from_pubsub(
    root,
    *,
    topic: Optional[str] = None,
    subscription: Optional[str] = None,
    format: str,
    schema: Optional[Any] = None,
    attributes: Optional[Iterable[str]] = None,
    attributes_map: Optional[str] = None,
    id_attribute: Optional[str] = None,
    timestamp_attribute: Optional[str] = None,
    publish_time_field: Optional[str] = None):
  """Reads messages from Cloud Pub/Sub.

  Args:
    topic: Cloud Pub/Sub topic in the form
      "projects/<project>/topics/<topic>". If provided, subscription must be
      None.
    subscription: Existing Cloud Pub/Sub subscription to use in the
      form "projects/<project>/subscriptions/<subscription>". If not
      specified, a temporary subscription will be created from the specified
      topic. If provided, topic must be None.
    format: The expected format of the message payload.  Currently suported
      formats are

        - RAW: Produces records with a single `payload` field whose contents
            are the raw bytes of the pubsub message.
        - STRING: Like RAW, but the bytes are decoded as a UTF-8 string.
        - AVRO: Parses records with a given Avro schema.
        - JSON: Parses records with a given JSON schema.

    schema: Schema specification for the given format.
    attributes: List of attribute keys whose values will be flattened into the
      output message as additional fields.  For example, if the format is `raw`
      and attributes is `["a", "b"]` then this read will produce elements of
      the form `Row(payload=..., a=..., b=...)`.
    attribute_map: Name of a field in which to store the full set of attributes
      associated with this message.  For example, if the format is `raw` and
      `attribute_map` is set to `"attrs"` then this read will produce elements
      of the form `Row(payload=..., attrs=...)` where `attrs` is a Map type
      of string to string.
      If both `attributes` and `attribute_map` are set, the overlapping
      attribute values will be present in both the flattened structure and the
      attribute map.
    id_attribute: The attribute on incoming Pub/Sub messages to use as a unique
      record identifier. When specified, the value of this attribute (which
      can be any string that uniquely identifies the record) will be used for
      deduplication of messages. If not provided, we cannot guarantee
      that no duplicate data will be delivered on the Pub/Sub stream. In this
      case, deduplication of the stream will be strictly best effort.
    timestamp_attribute: Message value to use as element timestamp. If None,
      uses message publishing time as the timestamp.

      Timestamp values should be in one of two formats:

      - A numerical value representing the number of milliseconds since the
        Unix epoch.
      - A string in RFC 3339 format, UTC timezone. Example:
        ``2015-10-29T23:41:41.123Z``. The sub-second component of the
        timestamp is optional, and digits beyond the first three (i.e., time
        units smaller than milliseconds) may be ignored.
    publish_time_field: Field to add to output messages with the Pub/Sub
      message publish time. If None, no such field is added.
  """
  if topic and subscription:
    raise TypeError('Only one of topic and subscription may be specified.')
  elif not topic and not subscription:
    raise TypeError('One of topic or subscription may be specified.')
  if publish_time_field is not None and not publish_time_field.strip():
    raise ValueError('publish_time_field must be a non-empty field name.')
  has_publish_time_field = publish_time_field is not None
  payload_schema, parser = _create_parser(format, schema)
  extra_fields: list[schema_pb2.Field] = []
  if not attributes and not attributes_map and not has_publish_time_field:
    mapper = lambda msg: parser(msg)
  else:
    if isinstance(attributes, str):
      attributes = [attributes]
    if attributes:
      extra_fields.extend(
          [schemas.schema_field(attr, str) for attr in attributes])
    if attributes_map:
      extra_fields.append(
          schemas.schema_field(attributes_map, Mapping[str, str]))
    if has_publish_time_field:
      extra_fields.append(
          schemas.schema_field(publish_time_field, Optional[Timestamp]))

    def mapper(msg):
      values = parser(msg.data).as_dict()
      if attributes:
        # Should missing attributes be optional or parse errors?
        for attr in attributes:
          values[attr] = msg.attributes[attr]
      if attributes_map:
        values[attributes_map] = msg.attributes
      if has_publish_time_field:
        values[publish_time_field] = (
            Timestamp.of(msg.publish_time)
            if msg.publish_time is not None else None)
      return beam.Row(**values)

  output = (
      root
      | beam.io.ReadFromPubSub(
          topic=topic,
          subscription=subscription,
          with_attributes=bool(
              attributes or attributes_map or has_publish_time_field),
          id_label=id_attribute,
          timestamp_attribute=timestamp_attribute)
      | 'ParseMessage' >> beam.Map(mapper))
  output.element_type = schemas.named_tuple_from_schema(
      schema_pb2.Schema(fields=list(payload_schema.fields) + extra_fields))
  return output


@beam.ptransform_fn
@yaml_errors.maybe_with_exception_handling_transform_fn
def write_to_pubsub(
    pcoll,
    *,
    topic: str,
    format: str,
    schema: Optional[Any] = None,
    attributes: Optional[Iterable[str]] = None,
    attributes_map: Optional[str] = None,
    id_attribute: Optional[str] = None,
    timestamp_attribute: Optional[str] = None):
  """Writes messages to Cloud Pub/Sub.

  Args:
    topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
    format: How to format the message payload.  Currently supported
      formats are

        - RAW: Expects a message with a single field (excluding
            attribute-related fields) whose contents are used as the raw bytes
            of the pubsub message.
        - AVRO: Encodes records with a given Avro schema, which may be inferred
            from the input PCollection schema.
        - JSON: Formats records with a given JSON schema, which may be inferred
            from the input PCollection schema.
        - PROTO: Encodes records with a given Protobuf schema, which may be
            inferred from the input PCollection schema.

    schema: Schema specification for the given format.
    attributes: List of attribute keys whose values will be pulled out as
      PubSub message attributes.  For example, if the format is `raw`
      and attributes is `["a", "b"]` then elements of the form
      `Row(any_field=..., a=..., b=...)` will result in PubSub messages whose
      payload has the contents of any_field and whose attribute will be
      populated with the values of `a` and `b`.
    attribute_map: Name of a string-to-string map field in which to pull a set
      of attributes associated with this message.  For example, if the format
      is `raw` and `attribute_map` is set to `"attrs"` then elements of the form
      `Row(any_field=..., attrs=...)` will result in PubSub messages whose
      payload has the contents of any_field and whose attribute will be
      populated with the values from attrs.
      If both `attributes` and `attribute_map` are set, the union of attributes
      from these two sources will be used to populate the PubSub message
      attributes.
    id_attribute: If set, will set an attribute for each Cloud Pub/Sub message
      with the given name and a unique value. This attribute can then be used
      in a ReadFromPubSub PTransform to deduplicate messages.
    timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub
      message with the given name and the message's publish time as the value.
  """
  input_schema = schemas.schema_from_element_type(pcoll.element_type)

  extra_fields: list[str] = []
  if isinstance(attributes, str):
    attributes = [attributes]
  if attributes:
    extra_fields.extend(attributes)
  if attributes_map:
    extra_fields.append(attributes_map)

  def attributes_extractor(row):
    if attributes_map:
      attribute_values = dict(getattr(row, attributes_map))
    else:
      attribute_values = {}
    if attributes:
      attribute_values.update({attr: getattr(row, attr) for attr in attributes})
    return attribute_values

  schema_names = set(f.name for f in input_schema.fields)
  missing_attribute_names = set(extra_fields) - schema_names
  if missing_attribute_names:
    raise ValueError(
        f'Attribute fields {missing_attribute_names} '
        f'not found in schema fields {schema_names}')

  payload_schema = schema_pb2.Schema(
      fields=[
          field for field in input_schema.fields
          if field.name not in extra_fields
      ])
  formatter = _create_formatter(format, schema, payload_schema)
  return (
      pcoll | beam.Map(
          lambda row: beam.io.gcp.pubsub.PubsubMessage(
              formatter(row), attributes_extractor(row)))
      | beam.io.WriteToPubSub(
          topic,
          with_attributes=True,
          id_label=id_attribute,
          timestamp_attribute=timestamp_attribute))


def read_from_iceberg(
    table: str,
    filter: Optional[str] = None,
    keep: Optional[list[str]] = None,
    drop: Optional[list[str]] = None,
    catalog_name: Optional[str] = None,
    catalog_properties: Optional[Mapping[str, str]] = None,
    config_properties: Optional[Mapping[str, str]] = None,
):
  # TODO(robertwb): It'd be nice to derive this list of parameters, along with
  # their types and docs, programmatically from the iceberg (or managed)
  # schemas.

  """Reads an Apache Iceberg table.

  See also the [Apache Iceberg Beam documentation](
  https://cloud.google.com/dataflow/docs/guides/managed-io#iceberg).

  Args:
    table: The identifier of the Apache Iceberg table. Example: "db.table1".
    filter: SQL-like predicate to filter data at scan time.
      Example: "id > 5 AND status = 'ACTIVE'". Uses Apache Calcite syntax:
      https://calcite.apache.org/docs/reference.html
    keep: A subset of column names to read exclusively. If null or empty,
      all columns will be read.
    drop: A subset of column names to exclude from reading. If null or empty,
      all columns will be read.
    catalog_name: The name of the catalog. Example: "local".
    catalog_properties: A map of configuration properties for the Apache Iceberg
      catalog.
      The required properties depend on the catalog. For more information, see
      CatalogUtil in the Apache Iceberg documentation.
    config_properties: An optional set of Hadoop configuration properties.
      For more information, see CatalogUtil in the Apache Iceberg documentation.
  """
  return beam.managed.Read(
      "iceberg",
      config=dict(
          table=table,
          catalog_name=catalog_name,
          filter=filter,
          keep=keep,
          drop=drop,
          catalog_properties=catalog_properties,
          config_properties=config_properties))


def read_from_delta(
    table: str,
    version: Optional[int] = None,
    timestamp: Optional[str] = None,
    hadoop_config: Optional[Mapping[str, str]] = None,
):
  """Reads a Delta Lake table.

  Args:
    table: Identifier of the Delta Lake table.
    version: Version of the Delta Lake table to read.
    timestamp: Timestamp of the Delta Lake table to read.
    hadoop_config: Hadoop configuration properties.
  """
  return beam.managed.Read(
      "delta",
      config=dict(
          table=table,
          version=version,
          timestamp=timestamp,
          hadoop_config=hadoop_config))


def read_from_delta_cdc(
    table: str,
    start_version: Optional[int] = None,
    start_timestamp: Optional[str] = None,
    end_version: Optional[int] = None,
    end_timestamp: Optional[str] = None,
    include_metadata_columns: Optional[Iterable[str]] = None,
    hadoop_config: Optional[Mapping[str, str]] = None,
):
  """Reads change records from a Delta Lake table.

  Args:
    table: Identifier of the Delta Lake table.
    start_version: Start version of the Delta Lake table to read changes from.
      Either this or start_timestamp has to be provided.
    start_timestamp: Start timestamp of the Delta Lake table to read changes
      from. Should be specified in the ISO 8601 standard. Either this or
      start_version has to be provided.
    end_version: End version of the Delta Lake table to read changes up to.
    end_timestamp: End timestamp of the Delta Lake table to read changes up to.
      Should be specified in the ISO 8601 standard.
    include_metadata_columns: Metadata columns to include in the output rows.
      Supported columns are: _change_type, _commit_version, and _commit_timestamp.
    hadoop_config: Properties passed to the Hadoop Configuration.
  """
  return beam.managed.Read(
      "delta_cdc",
      config=dict(
          table=table,
          start_version=start_version,
          start_timestamp=start_timestamp,
          end_version=end_version,
          end_timestamp=end_timestamp,
          include_metadata_columns=include_metadata_columns,
          hadoop_config=hadoop_config))


def write_to_iceberg(
    table: str,
    catalog_name: Optional[str] = None,
    catalog_properties: Optional[Mapping[str, str]] = None,
    config_properties: Optional[Mapping[str, str]] = None,
    partition_fields: Optional[Iterable[str]] = None,
    table_properties: Optional[Mapping[str, str]] = None,
    triggering_frequency_seconds: Optional[int] = None,
    keep: Optional[Iterable[str]] = None,
    drop: Optional[Iterable[str]] = None,
    only: Optional[str] = None,
    distribution_mode: Optional[str] = None,
    autosharding: Optional[bool] = None,
    mode: Optional[str] = None,
    sequence_number_column: Optional[str] = None,
    change_type_column: Optional[str] = None,
    change_type_map: Optional[Mapping[str, str]] = None,
    upsert: Optional[bool] = None,
    equality_columns: Optional[Iterable[str]] = None,
    num_shards: Optional[int] = None,
    shards_per_partition: Optional[int] = None,
    allowed_lateness_seconds: Optional[int] = None,
    sink_id: Optional[str] = None,
    token_heartbeat_seconds: Optional[int] = None,
    snapshot_properties: Optional[Mapping[str, str]] = None,
    error_handling: Optional[Mapping[str, Any]] = None,
    sorter_memory_mb: Optional[int] = None,
):
  # TODO(robertwb): It'd be nice to derive this list of parameters, along with
  # their types and docs, programmatically from the iceberg (or managed)
  # schemas.

  """Writes to an Apache Iceberg table.

  See also the [Apache Iceberg Beam documentation](
  https://cloud.google.com/dataflow/docs/guides/managed-io#iceberg)
  including the [dynamic destinations section](
  https://cloud.google.com/dataflow/docs/guides/managed-io#dynamic-destinations)
  for use of the keep, drop, and only parameters.

  Args:
    table: The identifier of the Apache Iceberg table. Example: "db.table1".
    catalog_name: The name of the catalog. Example: "local".
    catalog_properties: A map of configuration properties for the Apache Iceberg
      catalog.
      The required properties depend on the catalog. For more information, see
      CatalogUtil in the Apache Iceberg documentation.
    config_properties: An optional set of Hadoop configuration properties.
      For more information, see CatalogUtil in the Apache Iceberg documentation.
    partition_fields: Fields used to create a partition spec that is applied
      when tables are created. For a field 'foo', the available partition
      transforms are:

        - foo
        - truncate(foo, N)
        - bucket(foo, N)
        - hour(foo)
        - day(foo)
        - month(foo)
        - year(foo)
        - void(foo)
      For more information on partition transforms, please visit
      https://iceberg.apache.org/spec/#partition-transforms.
    table_properties: Iceberg table properties to be set on the table when it
      is created. For more information on table properties, please visit
      https://iceberg.apache.org/docs/latest/configuration/#table-properties.
    triggering_frequency_seconds: For streaming write pipelines, the frequency
      at which the sink attempts to produce snapshots, in seconds.
    keep: An optional list of field names to keep when writing to the
      destination. Other fields are dropped. Mutually exclusive with drop
      and only.
    drop: An optional list of field names to drop before writing to the
        destination. Mutually exclusive with keep and only.
    only: The name of exactly one field to keep as the top level record when
      writing to the destination. All other fields are dropped. This field must
      be of row type. Mutually exclusive with drop and keep.
    distribution_mode: Defines distribution of write data. Supported
      distributions:
      - none: don't shuffle rows (default)
      - hash: shuffle rows by partition key before writing data
    autosharding: Enables dynamic sharding to automatically adjust the number
      of parallel writers based on data volume. It handles data skew by
      further sub-dividing partitions into multiple shards to prevent
      bottlenecks during high-throughput writes. Only available with 'hash'
      distribution mode.
    mode: Controls how rows are written. 'append' (default) appends every row
      as new data. 'merge-on-read' treats each row as a change (INSERT,
      UPDATE_BEFORE, UPDATE_AFTER, or DELETE) applied to the table by primary
      key.
    sequence_number_column: Merge-on-read only. the required column name
      representing the monotonic sequence number used to order a single key's
      changes. Default column name is '_commit_snapshot_sequence_number'. This
      column will be stripped from the data row before writing to Iceberg.
    change_type_column: Merge-on-read only. The optional column name
      representing the row's change type (INSERT, UPDATE_BEFORE, UPDATE_AFTER,
      or DELETE). This column will be stripped from the data row before writing
      to Iceberg.
    change_type_map: Merge-on-read only. Optional map from the
      `change_type_column` value to the canonical change type name (see above).
    upsert: Merge-on-read only. If true, only the after-image of each change
      (INSERT/UPDATE_AFTER) is applied as an upsert. UPDATE_BEFORE records
      are dropped. Default: false.
    equality_columns: Columns defining row identity (equality-delete fields).
      Defaults to the destination table's identifier (primary-key) fields.
      Required if the table doesn't exist yet. Currently only supported in
      'merge-on-read' mode.
    num_shards: The number of deterministic primary-key-hash shards per
      destination, i.e. the max write parallelism per destination. Defaults to
      16. Currently only supported in 'merge-on-read' mode.
    shards_per_partition: Maximum number of shards a single partition's rows
      may occupy. Defaults to `num_shards`. Currently only supported in
      'merge-on-read' mode.
    allowed_lateness_seconds: How long a late record may lag behind the
      watermark before it is dropped entirely. Defaults to 21600 (6 hours).
      Currently only supported in 'merge-on-read' mode.
    sink_id: A stable identifier namespacing the idempotency tokens written to
      each commit. Defaults to a unique per-write UUID. Override with a stable
      ID for exactly-once commits across relaunches of a particular streaming
      write. Currently only relevant in 'merge-on-read' mode.
    token_heartbeat_seconds: Streaming only. Refresh each idle
      destination's commit token every this many seconds. Disabled by default.
      Currently only relevant in 'merge-on-read' mode.
    snapshot_properties: Extra properties added to every commit's snapshot
      summary. Currently only supported in 'merge-on-read' mode.
    error_handling: Where to send records the sink cannot apply (e.g. an
      unknown change type) instead of failing the pipeline. Currently only
      supported in 'merge-on-read' mode.
    sorter_memory_mb: The in-memory buffer size (MB) for the pre-write sort.
      Groups larger than this spill to disk. Default: 100MB. Currently only
      relevant in 'merge-on-read' mode.
  """
  return beam.managed.Write(
      "iceberg",
      config=dict(
          table=table,
          catalog_name=catalog_name,
          catalog_properties=catalog_properties,
          config_properties=config_properties,
          partition_fields=partition_fields,
          table_properties=table_properties,
          triggering_frequency_seconds=triggering_frequency_seconds,
          keep=keep,
          drop=drop,
          only=only,
          distribution_mode=distribution_mode,
          autosharding=autosharding,
          mode=mode,
          sequence_number_column=sequence_number_column,
          change_type_column=change_type_column,
          change_type_map=change_type_map,
          upsert=upsert,
          equality_columns=equality_columns,
          num_shards=num_shards,
          shards_per_partition=shards_per_partition,
          allowed_lateness_seconds=allowed_lateness_seconds,
          sink_id=sink_id,
          token_heartbeat_seconds=token_heartbeat_seconds,
          snapshot_properties=snapshot_properties,
          error_handling=error_handling,
          sorter_memory_mb=sorter_memory_mb))


def io_providers():
  return yaml_provider.load_providers(
      yaml_utils.locate_data_file('standard_io.yaml'))


def read_from_tfrecord(
    file_pattern: str,
    coder: Optional[coders.BytesCoder] = coders.BytesCoder(),
    compression_type: str = "AUTO",
    validate: Optional[bool] = True):
  """Reads data from TFRecord.

  Args:
    file_pattern (str): A file glob pattern to read TFRecords from.
    coder (coders.BytesCoder): Coder used to decode each record.
    compression_type (CompressionTypes): Used to handle compressed input files.
      Default value is CompressionTypes.AUTO, in which case the file_path's
      extension will be used to detect the compression.
    validate (bool): Boolean flag to verify that the files exist during the
      pipeline creation time.
  """
  return ReadFromTFRecord(
      file_pattern=file_pattern,
      compression_type=getattr(CompressionTypes, compression_type),
      validate=validate) | beam.Map(lambda s: beam.Row(record=s))


@beam.ptransform_fn
def write_to_tfrecord(
    pcoll,
    file_path_prefix: str,
    coder: Optional[coders.BytesCoder] = coders.BytesCoder(),
    file_name_suffix: Optional[str] = "",
    num_shards: Optional[int] = 0,
    shard_name_template: Optional[str] = None,
    compression_type: str = "AUTO"):
  """Writes data to TFRecord.

  Args:
    file_path_prefix: The file path to write to. The files written will begin
      with this prefix, followed by a shard identifier (see num_shards), and
      end in a common extension, if given by file_name_suffix.
    coder: Coder used to encode each record.
    file_name_suffix: Suffix for the files written.
    num_shards: The number of files (shards) used for output. If not set, the
      default value will be used.
    shard_name_template: A template string containing placeholders for
      the shard number and shard count. When constructing a filename for a
      particular shard number, the upper-case letters 'S' and 'N' are
      replaced with the 0-padded shard number and shard count respectively.
      This argument can be '' in which case it behaves as if num_shards was
      set to 1 and only one file will be generated. The default pattern used
      is '-SSSSS-of-NNNNN' if None is passed as the shard_name_template.
    compression_type: Used to handle compressed output files. Typical value
      is CompressionTypes.AUTO, in which case the file_path's extension will
      be used to detect the compression.

  Returns:
    A WriteToTFRecord transform object.
  """
  try:
    field_names = [
        name for name, _ in schemas.named_fields_from_element_type(
            pcoll.element_type)
    ]
  except Exception as exn:
    raise ValueError(
        "WriteToTFRecord requires an input schema with exactly one field."
    ) from exn
  if len(field_names) != 1:
    raise ValueError(
        "WriteToTFRecord requires an input schema with exactly one field,got %s"
        % field_names)
  sole_field_name, = field_names

  return pcoll | beam.Map(
      lambda x: getattr(x, sole_field_name)) | WriteToTFRecord(
          file_path_prefix=file_path_prefix,
          coder=coder,
          file_name_suffix=file_name_suffix,
          num_shards=num_shards,
          shard_name_template=shard_name_template,
          compression_type=getattr(CompressionTypes, compression_type))


@beam.ptransform_fn
@yaml_errors.maybe_with_exception_handling_transform_fn
def read_from_mongodb(
    root,
    *,
    database: str,
    collection: str,
    schema: Union[str, dict[str, Any]],
    uri: Optional[str] = None,
    filter: Optional[dict[str, Any]] = None):
  """Reads data from MongoDB.

  The resulting PCollection consists of rows with fields matching the provided
  schema.

  Args:
    database: The MongoDB database name.
    collection: The MongoDB collection name.
    schema: JSON schema specifying the fields to select and their types.
    uri: The MongoDB connection string. e.g. "mongodb://localhost:27017"
    filter: A JSON/bson mapping specifying elements which must be present.
  """
  if isinstance(schema, str):
    schema = json.loads(schema)
  if isinstance(filter, str):
    filter = json.loads(filter)

  beam_schema = json_utils.json_schema_to_beam_schema(schema)
  beam_type = schema_pb2.FieldType(
      row_type=schema_pb2.RowType(schema=beam_schema))
  to_row_fn = json_utils.json_to_row(beam_type)

  output = (
      root
      | mongodbio.ReadFromMongoDB(
          uri=uri, db=database, coll=collection, filter=filter)
      | beam.Map(to_row_fn))
  output.element_type = schemas.named_tuple_from_schema(beam_schema)
  return output


@beam.ptransform_fn
@yaml_errors.maybe_with_exception_handling_transform_fn
def write_to_mongodb(
    pcoll,
    *,
    database: str,
    collection: str,
    uri: Optional[str] = None,
    batch_size: int = 1024):
  """Writes data to MongoDB.

  Args:
    pcoll: The input PCollection of Beam Rows.
    database: The MongoDB database name.
    collection: The MongoDB collection name.
    uri: The MongoDB connection string. e.g. "mongodb://localhost:27017"
    batch_size: Number of documents per bulk_write to MongoDB.
  """
  def row_to_dict(value):
    if value is None:
      return None
    if hasattr(value, '_asdict'):
      return {k: row_to_dict(v) for k, v in value._asdict().items()}
    elif hasattr(value, 'as_dict'):
      return {k: row_to_dict(v) for k, v in value.as_dict().items()}
    elif isinstance(value, (list, tuple)):
      return [row_to_dict(v) for v in value]
    elif isinstance(value, Mapping):
      return {k: row_to_dict(v) for k, v in value.items()}
    else:
      return value

  return (
      pcoll
      | beam.Map(row_to_dict)
      | mongodbio.WriteToMongoDB(
          uri=uri, db=database, coll=collection, batch_size=batch_size))


@beam.ptransform_fn
def match_all(
    pcoll,
    *,
    file_pattern: Optional[str] = None,
    empty_match_treatment: str = 'ALLOW',
):
  """Matches file patterns from the input PCollection.

  This transform returns a PCollection of matching files, each represented as a
  Row with path, size_in_bytes, and last_updated_in_seconds fields.

  Args:
    file_pattern (str): The name of the field in the input PCollection that contains
      the file pattern string. If not specified and the input PCollection has
      exactly one field, that field will be used.
    empty_match_treatment (str): How to treat empty matches. Possible values are
      'ALLOW', 'DISALLOW', and 'ALLOW_IF_WILDCARD'. Defaults to 'ALLOW'.
  """
  from apache_beam.typehints import schemas

  try:
    field_names = [
        name for name, _ in schemas.named_fields_from_element_type(
            pcoll.element_type)
    ]
  except Exception:
    field_names = None

  if field_names:
    if file_pattern is not None:
      if file_pattern not in field_names:
        raise ValueError(
            f"Field '{file_pattern}' not found in input schema fields: {field_names}"
        )
      pattern_field = file_pattern
    elif len(field_names) == 1:
      pattern_field = field_names[0]
    else:
      raise ValueError(
          f"Input schema has multiple fields {field_names}. "
          f"Please specify the 'file_pattern' parameter to select which field "
          f"contains the file pattern.")
    patterns = pcoll | beam.Map(lambda x: str(getattr(x, pattern_field)))
  else:
    patterns = pcoll

  matched = patterns | beam.io.fileio.MatchAll(
      empty_match_treatment=empty_match_treatment)

  return matched | beam.Map(
      lambda x: beam.Row(
          path=str(x.path), size_in_bytes=int(x.size_in_bytes),
          last_updated_in_seconds=float(x.last_updated_in_seconds)
          if x.last_updated_in_seconds is not None else None))


_DICOM_SEARCH_OUTPUT_SCHEMA = RowTypeConstraint.from_fields([
    ('result', str),
    ('status', str),
    ('input', str),
])


def _dicom_search_result_to_row(result):
  if not result.get('success'):
    raise RuntimeError(
        'DicomSearch failed with status: %s' % (result.get('status'), ))
  return beam.Row(
      result=json.dumps(result.get('result', [])),
      status=str(result.get('status')),
      input=json.dumps(result.get('input', {})))


def _dicom_search_to_output_row(result):
  return beam.Row(
      result=json.dumps(result.get('result', [])),
      status=str(result.get('status')),
      input=json.dumps(result.get('input', {})))


def _dicom_search_to_error_row(result):
  inp = result.get('input') or {}
  if isinstance(inp, Mapping):
    element = beam.Row(**dict(inp))
  else:
    element = inp
  return beam.Row(
      element=element,
      msg='DicomSearch failed with status: %s' % (result.get('status'), ),
      stack='')


@beam.ptransform_fn
def dicom_search(
    pcoll, *, buffer_size: int = 8, max_workers: int = 5, error_handling=None):
  """Searches a Google Cloud Healthcare DICOM store using QIDO-RS.

  This transform takes an input PCollection of Rows describing QIDO search
  requests and returns Rows with the search results encoded as JSON.

  Each input Row must include:

    - project_id (str): GCP project containing the DICOM store.
    - region (str): Region where the DICOM store resides.
    - dataset_id (str): Dataset containing the DICOM store.
    - dicom_store_id (str): DICOM store id.
    - search_type (str): One of ``studies``, ``series``, or ``instances``.
    - params (map of str to str, optional): QIDO search filters.

  Successful outputs are Rows with:

    - result (str): JSON-encoded list of matching DICOM resources.
    - status (str): HTTP status from the DICOM API.
    - input (str): JSON-encoded copy of the search request.

  Failed searches raise unless ``error_handling`` is set, in which case they
  are routed to the configured error output.

  Args:
    buffer_size: Number of requests to buffer before flushing.
    max_workers: Maximum number of threads used to issue requests.
    error_handling: If specified, should be a mapping giving an output into
      which to emit failed searches, as described at
      https://beam.apache.org/documentation/sdks/yaml-errors/
  """
  try:
    from apache_beam.io.gcp.healthcare.dicomio import DicomSearch
  except ImportError as exn:
    raise ValueError(
        "GCP dependencies are not installed. Cannot use DicomSearch. "
        "Please install using 'pip install apache-beam[gcp]'.") from exn

  def row_to_dict(value):
    if value is None:
      return None
    if hasattr(value, '_asdict'):
      return {k: row_to_dict(v) for k, v in value._asdict().items()}
    elif hasattr(value, 'as_dict'):
      return {k: row_to_dict(v) for k, v in value.as_dict().items()}
    elif isinstance(value, (list, tuple)):
      return [row_to_dict(v) for v in value]
    elif isinstance(value, Mapping):
      return {k: row_to_dict(v) for k, v in value.items()}
    else:
      return value

  def normalize_request(value):
    # YAML Create types params as map[str, str]; qido_search needs int
    # limit/offset for pagination comparisons.
    request = row_to_dict(value)
    params = request.get('params')
    params = dict(params) if isinstance(params, Mapping) else {}
    limit = params.get('limit', 500)
    offset = params.get('offset', 0)
    params['limit'] = int(limit)
    params['offset'] = int(offset)
    request['params'] = params
    return request

  if error_handling:
    error_handling = yaml_utils.SafeLineLoader.strip_metadata(error_handling)

  results = (
      pcoll
      | beam.Map(normalize_request)
      | DicomSearch(buffer_size=buffer_size, max_workers=max_workers))

  if error_handling and error_handling.get('output'):
    return {
        'good': (
            results
            | beam.Filter(lambda r: r.get('success'))
            | beam.Map(_dicom_search_to_output_row).with_output_types(
                _DICOM_SEARCH_OUTPUT_SCHEMA)),
        error_handling['output']: (
            results
            | beam.Filter(lambda r: not r.get('success'))
            | beam.Map(_dicom_search_to_error_row)),
    }

  return results | beam.Map(_dicom_search_result_to_row).with_output_types(
      _DICOM_SEARCH_OUTPUT_SCHEMA)