sdks/python/apache_beam/yaml/yaml_mapping.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 defines the basic MapToFields operation."""
import datetime
import itertools
import json
import re
import threading
from collections import abc
from collections.abc import Callable
from collections.abc import Collection
from collections.abc import Iterable
from collections.abc import Mapping
from decimal import Decimal
from typing import Any
from typing import NamedTuple
from typing import Optional
from typing import TypeVar
from typing import Union

import apache_beam as beam
from apache_beam.io.filesystems import FileSystems
from apache_beam.portability.api import schema_pb2
from apache_beam.transforms.window import TimestampedValue
from apache_beam.typehints import row_type
from apache_beam.typehints import schemas
from apache_beam.typehints import trivial_inference
from apache_beam.typehints import typehints
from apache_beam.typehints.native_type_compatibility import convert_to_beam_type
from apache_beam.typehints.schemas import named_fields_from_element_type
from apache_beam.typehints.schemas import schema_from_element_type
from apache_beam.typehints.schemas import typing_from_runner_api
from apache_beam.utils import python_callable
from apache_beam.utils import windowed_value
from apache_beam.utils.timestamp import Timestamp
from apache_beam.yaml import json_utils
from apache_beam.yaml import options
from apache_beam.yaml import yaml_provider
from apache_beam.yaml.yaml_errors import exception_handling_args
from apache_beam.yaml.yaml_errors import map_errors_to_standard_format
from apache_beam.yaml.yaml_errors import maybe_with_exception_handling
from apache_beam.yaml.yaml_errors import maybe_with_exception_handling_transform_fn
from apache_beam.yaml.yaml_provider import dicts_to_rows

# Import quickjs package if it exists
try:
  import quickjs
except ImportError:
  quickjs = None

_str_expression_fields = {
    'AssignTimestamps': 'timestamp',
    'Filter': 'keep',
    'Partition': 'by',
}


def normalize_mapping(spec):
  """
  Normalizes various fields for mapping transforms.
  """
  if spec['type'] == 'MapToFields':
    config = spec.get('config')
    if isinstance(config.get('drop'), str):
      config['drop'] = [config['drop']]
    for field, value in list(config.get('fields', {}).items()):
      if isinstance(value, (str, int, float)):
        config['fields'][field] = {'expression': str(value)}

  elif spec['type'] in _str_expression_fields:
    param = _str_expression_fields[spec['type']]
    config = spec.get('config', {})
    if isinstance(config.get(param), (str, int, float)):
      config[param] = {'expression': str(config.get(param))}

  return spec


def is_literal(expr: str) -> bool:
  # Some languages have limited integer literal ranges.
  if re.fullmatch(r'-?\d+?', expr) and -1 << 31 < int(expr) < 1 << 31:
    return True
  elif re.fullmatch(r'-?\d+\.\d*', expr):
    return True
  elif re.fullmatch(r'"[^\\"]*"', expr):
    return True
  else:
    return False


def validate_generic_expression(
    expr_dict: dict,
    input_fields: Collection[str],
    allow_cmp: bool,
    error_field: str) -> None:
  if not isinstance(expr_dict, dict):
    raise ValueError(
        f"Ambiguous expression type (perhaps missing quoting?): {expr_dict}")
  if len(expr_dict) != 1 or 'expression' not in expr_dict:
    raise ValueError(
        "Missing language specification. "
        "Must specify a language when using a map with custom logic for %s" %
        error_field)
  expr = str(expr_dict['expression'])

  def is_atomic(expr: str):
    return is_literal(expr) or expr in input_fields

  if is_atomic(expr):
    return

  if allow_cmp:
    maybe_cmp = re.fullmatch('(.*)([<>=!]+)(.*)', expr)
    if maybe_cmp:
      left, cmp, right = maybe_cmp.groups()
      if (is_atomic(left.strip()) and is_atomic(right.strip()) and
          cmp in {'==', '<=', '>=', '<', '>', '!='}):
        return

  raise ValueError(
      "Missing language specification, unknown input fields, "
      f"or invalid generic expression: {expr}. "
      f"The given input fields are {input_fields}. "
      "See https://beam.apache.org/documentation/sdks/yaml-udf/#generic")


def validate_generic_expressions(base_type, config, input_pcolls) -> None:
  if not input_pcolls:
    return
  try:
    input_fields = [
        name for (name, _) in named_fields_from_element_type(
            next(iter(input_pcolls)).element_type)
    ]
  except (TypeError, ValueError):
    input_fields = []

  if base_type == 'MapToFields':
    for field, value in list(config.get('fields', {}).items()):
      validate_generic_expression(value, input_fields, True, field)

  elif base_type in _str_expression_fields:
    param = _str_expression_fields[base_type]
    validate_generic_expression(
        config.get(param), input_fields, base_type == 'Filter', param)


def _check_mapping_arguments(
    transform_name, expression=None, callable=None, name=None, path=None):
  # Argument checking
  if not expression and not callable and not path and not name:
    raise ValueError(
        f'{transform_name} must specify either "expression", "callable", '
        f'or both "path" and "name"')
  if expression and callable:
    raise ValueError(
        f'{transform_name} cannot specify both "expression" and "callable"')
  if (expression or callable) and (path or name):
    raise ValueError(
        f'{transform_name} cannot specify "expression" or "callable" with '
        f'"path" or "name"')
  if path and not name:
    raise ValueError(f'{transform_name} cannot specify "path" without "name"')
  if name and not path:
    raise ValueError(f'{transform_name} cannot specify "name" without "path"')


_THREAD_LOCAL_JS_CACHE = threading.local()


class _JsFunctionWrapper:
  def __init__(self, source_code, entrypoint_name):
    self.source_code = source_code
    self.entrypoint_name = entrypoint_name

  def _get_fn(self):
    cache = _THREAD_LOCAL_JS_CACHE
    if not hasattr(cache, 'functions'):
      cache.functions = {}

    cache_key = (self.source_code, self.entrypoint_name)
    if cache_key not in cache.functions:
      context = quickjs.Context()
      context.eval(self.source_code)
      f = context.get(self.entrypoint_name)

      def call_fn(*args, run_gc=True):
        def convert_arg(arg):
          if isinstance(arg, (type(None), str, bool, float, int)):
            return arg
          else:
            return context.parse_json(json.dumps(arg))

        try:
          result = f(*[convert_arg(a) for a in args])
          if isinstance(result, quickjs.Object):
            result = json.loads(result.json())
          return result
        finally:
          if run_gc:
            context.gc()

      cache.functions[cache_key] = call_fn

    return cache.functions[cache_key]

  def __call__(self, row):
    fn = self._get_fn()
    try:
      return dicts_to_rows(fn(py_value_to_js_dict(row)))
    except Exception as exn:
      raise RuntimeError(
          f"Error evaluating javascript expression: {exn}") from exn


# TODO(yaml) Improve type inferencing for JS UDF's
def py_value_to_js_dict(py_value):
  if ((isinstance(py_value, tuple) and hasattr(py_value, '_asdict')) or
      isinstance(py_value, beam.Row)):
    py_value = py_value._asdict()
  if isinstance(py_value, dict):
    return {key: py_value_to_js_dict(value) for key, value in py_value.items()}
  elif isinstance(py_value, bytes):
    return py_value.decode('utf-8', errors='replace')
  elif isinstance(py_value, (datetime.datetime, datetime.date, datetime.time)):
    return {'__date__': True, 'value': py_value.isoformat()}
  elif isinstance(py_value, Decimal):
    return float(py_value)
  elif not isinstance(py_value, str) and isinstance(py_value, abc.Iterable):
    return [py_value_to_js_dict(value) for value in list(py_value)]
  else:
    return py_value


# TODO(yaml) Consider adding optional language version parameter to support
#  ECMAScript 5 and 6
def _expand_javascript_mapping_func(
    original_fields, expression=None, callable=None, path=None, name=None):

  if quickjs is None:
    raise ValueError(
        "Javascript mapping functions are not supported because the "
        "quickjs-ng library is not installed.")

  if expression:
    source_code = f"""
    function udf(__row__) {{
      with (__row__) {{
        return ({expression});
      }}
    }}
    """
    user_entrypoint = 'udf'

  elif callable:
    source_code = f"var __udf__ = ({callable});"
    user_entrypoint = '__udf__'

  else:
    if not path.endswith('.js'):
      raise ValueError(f'File "{path}" is not a valid .js file.')
    udf_code = FileSystems.open(path).read().decode()
    source_code = udf_code
    user_entrypoint = name

  source_code += f"""
  function __convert_dates__(obj) {{
    if (obj && typeof obj === 'object') {{
      if (obj.__date__) {{
        return new Date(obj.value);
      }}
      for (var key in obj) {{
        if (obj.hasOwnProperty(key)) {{
          obj[key] = __convert_dates__(obj[key]);
        }}
      }}
    }}
    return obj;
  }}

  function __wrapper__(row) {{
    return {user_entrypoint}(__convert_dates__(row));
  }}
  """

  return _JsFunctionWrapper(source_code, '__wrapper__')


def _expand_python_mapping_func(
    original_fields, expression=None, callable=None, path=None, name=None):
  if path and name:
    if not path.endswith('.py'):
      raise ValueError(f'File "{path}" is not a valid .py file.')
    py_file = FileSystems.open(path).read().decode()

    return python_callable.PythonCallableWithSource.load_from_script(
        py_file, name)

  elif expression:
    # TODO(robertwb): Consider constructing a single callable that takes
    # the row and returns the new row, rather than invoking (and unpacking)
    # for each field individually.
    source = '\n'.join(['def fn(__row__):'] + [
        f'  {name} = __row__.{name}'
        for name in original_fields if name in expression
    ] + ['  return (' + expression + ')'])

  else:
    source = callable

  return python_callable.PythonCallableWithSource(source)


def _validator(beam_type: schema_pb2.FieldType) -> Callable[[Any], bool]:
  """Returns a callable converting rows of the given type to Json objects."""
  type_info = beam_type.WhichOneof("type_info")
  if type_info == "atomic_type":
    if beam_type.atomic_type == schema_pb2.BOOLEAN:
      return lambda x: isinstance(x, bool)
    elif beam_type.atomic_type == schema_pb2.INT64:
      return lambda x: isinstance(x, int)
    elif beam_type.atomic_type == schema_pb2.DOUBLE:
      return lambda x: isinstance(x, (int, float))
    elif beam_type.atomic_type == schema_pb2.STRING:
      return lambda x: isinstance(x, str)
    elif beam_type.atomic_type == schema_pb2.BYTES:
      return lambda x: isinstance(x, bytes)
    else:
      raise ValueError(
          f'Unknown or unsupported atomic type: {beam_type.atomic_type}')
  elif type_info == "array_type":
    element_validator = _validator(beam_type.array_type.element_type)
    return lambda value: all(element_validator(e) for e in value)
  elif type_info == "iterable_type":
    element_validator = _validator(beam_type.iterable_type.element_type)
    return lambda value: all(element_validator(e) for e in value)
  elif type_info == "map_type":
    key_validator = _validator(beam_type.map_type.key_type)
    value_validator = _validator(beam_type.map_type.value_type)
    return lambda value: all(
        key_validator(k) and value_validator(v) for (k, v) in value.items())
  elif type_info == "row_type":
    validators = {
        field.name: _validator(field.type)
        for field in beam_type.row_type.schema.fields
    }
    return lambda row: all(
        validator(getattr(row, name))
        for (name, validator) in validators.items())
  else:
    raise ValueError(f"Unrecognized type_info: {type_info!r}")


def _as_callable_for_pcoll(
    pcoll,
    fn_spec: Union[str, dict[str, str]],
    msg: str,
    language: Optional[str]):
  if language == 'javascript':
    options.YamlOptions.check_enabled(pcoll.pipeline, 'javascript')

  try:
    input_schema = dict(named_fields_from_element_type(pcoll.element_type))
  except (TypeError, ValueError) as exn:
    if is_expr(fn_spec):
      raise ValueError("Can only use expressions on a schema'd input.") from exn
    input_schema = {}  # unused

  if isinstance(fn_spec, str) and fn_spec in input_schema:
    return lambda row: getattr(row, fn_spec)
  else:
    return _as_callable(
        list(input_schema.keys()), fn_spec, msg, language, input_schema)


def _as_callable(original_fields, expr, transform_name, language, input_schema):
  if isinstance(expr, str):
    expr = {'expression': expr}

  # Extract original type from upstream pcoll when doing simple mappings
  original_type = input_schema.get(expr.get('expression'), None)
  if expr in original_fields:
    language = "python"

  # TODO(yaml): support an imports parameter
  # TODO(yaml): support a requirements parameter (possibly at a higher level)
  if not isinstance(expr, dict):
    raise ValueError(
        f"Ambiguous expression type (perhaps missing quoting?): {expr}")
  explicit_type = expr.pop('output_type', None)
  _check_mapping_arguments(transform_name, **expr)

  if language == "javascript":
    func = _expand_javascript_mapping_func(original_fields, **expr)
  elif language in ("python", "generic", None):
    func = _expand_python_mapping_func(original_fields, **expr)
  else:
    raise ValueError(
        f'Unknown language for mapping transform: {language}. '
        'Supported languages are "javascript" and "python."')

  if explicit_type:
    if isinstance(explicit_type, str):
      explicit_type = {'type': explicit_type}
    beam_type = json_utils.json_type_to_beam_type(explicit_type)
    validator = _validator(beam_type)

    @beam.typehints.with_output_types(schemas.typing_from_runner_api(beam_type))
    def checking_func(row):
      result = func(row)
      if not validator(result):
        raise TypeError(f'{result} violates schema {explicit_type}')
      return result

    return checking_func

  elif original_type:
    return beam.typehints.with_output_types(
        convert_to_beam_type(original_type))(
            func)

  else:
    return func


class _StripErrorMetadata(beam.PTransform):
  """Strips error metadata from outputs returned via error handling.

  Generally the error outputs for transformations return information about
  the error encountered (e.g. error messages and tracebacks) in addition to the
  failing element itself.  This transformation attempts to remove that metadata
  and returns the bad element alone which can be useful for re-processing.

  For example, in the following pipeline snippet::

      - name: MyMappingTransform
        type: MapToFields
        input: SomeInput
        config:
          language: python
          fields:
            ...
          error_handling:
            output: errors

      - name: RecoverOriginalElements
        type: StripErrorMetadata
        input: MyMappingTransform.errors

  the output of `RecoverOriginalElements` will contain exactly those elements
  from SomeInput that failed to processes (whereas `MyMappingTransform.errors`
  would contain those elements paired with error information).

  Note that this relies on the preceding transform actually returning the
  failing input in a schema'd way.  Most built-in transformation follow the
  correct conventions.
  """

  _ERROR_FIELD_NAMES = ('failed_row', 'element', 'record')

  def __init__(self):
    super().__init__(label=None)

  def expand(self, pcoll):
    try:
      existing_fields = {
          fld.name: fld.type
          for fld in schema_from_element_type(pcoll.element_type).fields
      }
    except TypeError:
      fld = None
    else:
      for fld in self._ERROR_FIELD_NAMES:
        if fld in existing_fields:
          break
      else:
        raise ValueError(
            'The input to this transform does not appear to be an error ' +
            "output.  Expected a schema'd input with a field named " +
            ' or '.join(repr(fld) for fld in self._ERROR_FIELD_NAMES))

    if fld is None:
      # This handles with_exception_handling() that returns bare tuples.
      return pcoll | beam.Map(lambda x: x[0])
    else:
      return pcoll | beam.Map(lambda x: getattr(x, fld)).with_output_types(
          typing_from_runner_api(existing_fields[fld]))


class Validate(beam.PTransform):
  """Validates each element of a PCollection against a json schema.

  Args:
      schema: A json schema against which to validate each element.
      error_handling: Whether and how to handle errors during iteration.
          If this is not set, invalid elements will fail the pipeline, otherwise
          invalid elements will be passed to the specified error output along
          with information about how the schema was invalidated.
  """
  def __init__(
      self,
      schema: dict[str, Any],
      error_handling: Optional[Mapping[str, Any]] = None):
    self._schema = schema
    self._exception_handling_args = exception_handling_args(error_handling)

  @maybe_with_exception_handling
  def expand(self, pcoll):
    validator = json_utils.row_validator(
        schema_from_element_type(pcoll.element_type), self._schema)

    def invoke_validator(x):
      validator(x)
      return x

    return pcoll | beam.Map(invoke_validator)

  def with_exception_handling(self, **kwargs):
    # It's possible there's an error in iteration...
    self._exception_handling_args = kwargs
    return self


class _Explode(beam.PTransform):
  """Explodes (aka unnest/flatten) one or more fields producing multiple rows.

  Given one or more fields of iterable type, produces multiple rows, one for
  each value of that field. For example, a row of the form `('a', [1, 2, 3])`
  would expand to `('a', 1)`, `('a', 2')`, and `('a', 3)` when exploded on
  the second field.

  This is akin to a `FlatMap` when paired with the MapToFields transform.

  See more complete documentation on
  [YAML Mapping Functions](https://beam.apache.org/documentation/sdks/yaml-udf/#flatmap).

  Args:
      fields: The list of fields to expand.
      cross_product: If multiple fields are specified, indicates whether the
          full cross-product of combinations should be produced, or if the
          first element of the first field corresponds to the first element
          of the second field, etc. For example, the row
          `(['a', 'b'], [1, 2])` would expand to the four rows
          `('a', 1)`, `('a', 2)`, `('b', 1)`, and `('b', 2)` when
          `cross_product` is set to `true` but only the two rows
          `('a', 1)` and `('b', 2)` when it is set to `false`.
          Only meaningful (and required) if multiple rows are specified.
      error_handling: Whether and how to handle errors during iteration.
  """  # pylint: disable=line-too-long

  def __init__(
      self,
      fields: Union[str, Collection[str]],
      cross_product: Optional[bool] = None,
      error_handling: Optional[Mapping[str, Any]] = None):
    if isinstance(fields, str):
      fields = [fields]
    if cross_product is None:
      if len(fields) > 1:
        raise ValueError(
            'cross_product must be specified true or false '
            'when exploding multiple fields')
      else:
        # Doesn't matter.
        cross_product = True
    self._fields = fields
    self._cross_product = cross_product
    # TODO(yaml):
    # 1. Support standard error handling argument.
    # 2. Supposedly error_handling parameter is not an accepted parameter when
    #    executing.  Needs further investigation.
    self._exception_handling_args = exception_handling_args(error_handling)

  @maybe_with_exception_handling
  def expand(self, pcoll):
    all_fields = [
        x for x, _ in named_fields_from_element_type(pcoll.element_type)
    ]
    for field in self._fields:
      if field not in all_fields:
        raise ValueError(f'Exploding unknown field "{field}"')
    to_explode = self._fields

    def explode_cross_product(base, fields):
      if fields:
        copy = dict(base)
        for value in base[fields[0]]:
          copy[fields[0]] = value
          yield from explode_cross_product(copy, fields[1:])
      else:
        yield beam.Row(**base)

    def explode_zip(base, fields):
      to_zip = [base[field] for field in fields]
      copy = dict(base)
      for values in itertools.zip_longest(*to_zip, fillvalue=None):
        for ix, field in enumerate(fields):
          copy[field] = values[ix]
        yield beam.Row(**copy)

    cross_product = self._cross_product
    return (
        pcoll
        | beam.FlatMap(
            lambda row:
            (explode_cross_product if cross_product else explode_zip)({
                name: getattr(row, name)
                for name in all_fields
            }, to_explode)))

  def infer_output_type(self, input_type):
    return row_type.RowTypeConstraint.from_fields([(
        name,
        trivial_inference.element_type(typ) if name in self._fields else
        typ) for (name, typ) in named_fields_from_element_type(input_type)])

  def with_exception_handling(self, **kwargs):
    # It's possible there's an error in iteration...
    self._exception_handling_args = kwargs
    return self


@beam.ptransform.ptransform_fn
@maybe_with_exception_handling_transform_fn
def _PyJsFilter(
    pcoll, keep: Union[str, dict[str, str]], language: Optional[str] = None):
  """Keeps only records that satisfy the given criteria.

  See more complete documentation on
  [YAML Filtering](https://beam.apache.org/documentation/sdks/yaml-udf/#filtering).

  Args:
    keep: An expression evaluating to true for those records that should be kept.
    language: The language of the above expression.
      Defaults to generic.
    error_handling: Whether and where to output records that throw errors when
      the above expressions are evaluated.
  """  # pylint: disable=line-too-long
  keep_fn = _as_callable_for_pcoll(pcoll, keep, "keep", language or 'generic')
  return pcoll | beam.Filter(keep_fn)


def is_expr(v):
  return isinstance(v, str) or (isinstance(v, dict) and 'expression' in v)


def normalize_fields(pcoll, fields, drop=(), append=False, language='generic'):
  try:
    input_schema = dict(named_fields_from_element_type(pcoll.element_type))
  except (TypeError, ValueError) as exn:
    if drop:
      raise ValueError("Can only drop fields on a schema'd input.") from exn
    if append:
      raise ValueError("Can only append fields on a schema'd input.") from exn
    elif any(is_expr(x) for x in fields.values()):
      raise ValueError("Can only use expressions on a schema'd input.") from exn
    input_schema = {}

  if drop and not append:
    raise ValueError("Can only drop fields if append is true.")
  for name in drop:
    if name not in input_schema:
      raise ValueError(f'Dropping unknown field "{name}"')
  if append:
    for name in fields:
      if name in input_schema and name not in drop:
        raise ValueError(
            f'Redefinition of field "{name}". '
            'Cannot append a field that already exists in original input.')

  if append:
    return input_schema, {
        **{name: f'`{name}`' if language in ['sql', 'calcite'] else name
           for name in input_schema.keys() if name not in drop},
        **fields
    }
  else:
    return input_schema, fields


@beam.ptransform.ptransform_fn
@maybe_with_exception_handling_transform_fn
def _PyJsMapToFields(
    pcoll,
    fields: Mapping[str, Union[str, Mapping[str, str]]],
    append: Optional[bool] = False,
    drop: Optional[Iterable[str]] = None,
    language: Optional[str] = None,
    dependencies: Optional[Iterable[str]] = None):
  """Creates records with new fields defined in terms of the input fields.

  See more complete documentation on
  [YAML Mapping Functions](https://beam.apache.org/documentation/sdks/yaml-udf/#mapping-functions).

  Args:
    fields: The output fields to compute, each mapping to the expression or
      callable that creates them.
    append: Whether to append the created fields to the set of
      fields already present, outputting a union of both the new fields and
      the original fields for each record.  Defaults to False.
    drop: If `append` is true, enumerates a subset of fields from the
      original record that should not be kept
    language: The language used to define (and execute) the
      expressions and/or callables in `fields`. Defaults to generic.
    dependencies: An optional list of extra dependencies that are needed for
      these UDFs.  The interpretation of these strings is language-dependent.
    error_handling: Whether and where to output records that throw errors when
      the above expressions are evaluated.
  """  # pylint: disable=line-too-long
  input_schema, fields = normalize_fields(
      pcoll, fields, drop or (), append, language=language or 'generic')
  if language == 'javascript':
    options.YamlOptions.check_enabled(pcoll.pipeline, 'javascript')

  original_fields = list(input_schema.keys())

  return pcoll | beam.Select(
      **{
          name: _as_callable(
              original_fields, expr, name, language, input_schema)
          for (name, expr) in fields.items()
      })


@beam.ptransform.ptransform_fn
def _SqlFilterTransform(pcoll, sql_transform_constructor, keep, language):
  return pcoll | sql_transform_constructor(
      f"SELECT * FROM PCOLLECTION WHERE {keep.get('expression')}")


@beam.ptransform.ptransform_fn
def _SqlMapToFieldsTransform(pcoll, sql_transform_constructor, **mapping_args):
  _, fields = normalize_fields(pcoll, **mapping_args)

  def extract_expr(name, v):
    if isinstance(v, str):
      return v
    elif 'expression' in v:
      return v['expression']
    else:
      raise ValueError(f"Only expressions allowed in SQL at {name}.")

  selects = [
      f'({extract_expr(name, expr)}) AS `{name}`'
      for (name, expr) in fields.items()
  ]
  query = "SELECT " + ", ".join(selects) + " FROM PCOLLECTION"
  return pcoll | sql_transform_constructor(query)


@beam.ptransform.ptransform_fn
def _Partition(
    pcoll,
    by: Union[str, dict[str, str]],
    outputs: list[str],
    unknown_output: Optional[str] = None,
    error_handling: Optional[Mapping[str, Any]] = None,
    language: str = 'generic'):
  """Splits an input into several distinct outputs.

  Each input element will go to a distinct output based on the field or
  function given in the `by` configuration parameter.

  Args:
      by: A field, callable, or expression giving the destination output for
        this element.  Should return a string that is a member of the `outputs`
        parameter. If `unknown_output` is also set, other returns values are
        accepted as well, otherwise an error will be raised.
      outputs: The set of outputs into which this input is being partitioned.
      unknown_output: (Optional) If set, indicates a destination output for any
        elements that are not assigned an output listed in the `outputs`
        parameter.
      error_handling: (Optional) Whether and how to handle errors during
        partitioning.
      language: The language of the `by` expression.
  """
  split_fn = _as_callable_for_pcoll(pcoll, by, 'by', language)
  try:
    split_fn_output_type = trivial_inference.infer_return_type(
        split_fn, [pcoll.element_type])
  except (TypeError, ValueError):
    pass
  else:
    if not typehints.is_consistent_with(split_fn_output_type,
                                        typehints.Optional[str]):
      raise ValueError(
          f'Partition function "{by}" must return a string type '
          f'not {split_fn_output_type}')
  error_output = error_handling['output'] if error_handling else None
  if error_output in outputs:
    raise ValueError(
        f'Error handling output "{error_output}" '
        f'cannot be among the listed outputs {outputs}')
  T = TypeVar('T')

  def split(element):
    tag = split_fn(element)
    if tag is None:
      tag = unknown_output
    if not isinstance(tag, str):
      raise ValueError(
          f'Returned output name "{tag}" of type {type(tag)} '
          f'from "{by}" must be a string.')
    if tag not in outputs:
      if unknown_output:
        tag = unknown_output
      else:
        raise ValueError(f'Unknown output name "{tag}" from {by}')
    return beam.pvalue.TaggedOutput(tag, element)

  output_set = set(outputs)
  if unknown_output:
    output_set.add(unknown_output)
  if error_output:
    output_set.add(error_output)
  mapping_transform = beam.Map(split)
  if error_output:
    mapping_transform = mapping_transform.with_exception_handling(
        **exception_handling_args(error_handling))
  else:
    mapping_transform = mapping_transform.with_outputs(*output_set)
  splits = pcoll | mapping_transform.with_input_types(T).with_output_types(T)
  result = {out: getattr(splits, out) for out in output_set}
  for tag, out in result.items():
    if tag != error_output:
      out.element_type = pcoll.element_type
  if error_output:
    result[error_output] = result[error_output] | map_errors_to_standard_format(
        pcoll.element_type)
  return result


@beam.ptransform.ptransform_fn
@maybe_with_exception_handling_transform_fn
def _AssignTimestamps(
    pcoll,
    timestamp: Union[str, dict[str, str]],
    language: Optional[str] = None):
  """Assigns a new timestamp each element of its input.

  This can be useful when reading records that have the timestamp embedded
  in them, for example with various file types or other sources that by default
  set all timestamps to the infinite past.

  Note that the timestamp should only be set forward, as setting it backwards
  may not cause it to hold back an already advanced watermark and the data
  could become droppably late.

  Args:
      timestamp: A field, callable, or expression giving the new timestamp.
      language: The language of the timestamp expression.
      error_handling: Whether and how to handle errors during timestamp
        evaluation.
  """
  timestamp_fn = _as_callable_for_pcoll(pcoll, timestamp, 'timestamp', language)
  T = TypeVar('T')
  return pcoll | beam.Map(lambda x: TimestampedValue(x, timestamp_fn(x))
                          ).with_input_types(T).with_output_types(T)


class PaneInfoTuple(NamedTuple):
  is_first: bool
  is_last: bool
  timing: str
  index: int  # type: ignore[assignment]
  nonspeculative_index: int

  @classmethod
  def from_pane_info(cls, pane_info):
    return cls(
        pane_info.is_first,
        pane_info.is_last,
        windowed_value.PaneInfoTiming.to_string(pane_info.timing),
        pane_info.index,
        pane_info.nonspeculative_index)


_WINDOWING_INFO_TYPES = {
    'timestamp': Timestamp,
    'window_start': Optional[Timestamp],
    'window_end': Timestamp,
    'window_string': str,
    'window_type': str,
    'window_object': Any,
    'pane_info': PaneInfoTuple,
}
_WINDOWING_INFO_EXTRACTORS = {
    'timestamp': lambda locals: locals['timestamp'],
    'window_start': lambda locals: getattr(locals['window'], 'start', None),
    'window_end': lambda locals: locals['window'].end,
    'window_string': lambda locals: str(locals['window']),
    'window_type': lambda locals: type(locals['window']).__name__,
    'window_object': lambda locals: locals['window'],
    'pane_info': lambda locals: PaneInfoTuple.from_pane_info(
        locals['pane_info']),
}
assert set(_WINDOWING_INFO_TYPES.keys()) == set(
    _WINDOWING_INFO_EXTRACTORS.keys())


@beam.ptransform.ptransform_fn
def _ExtractWindowingInfo(
    pcoll, fields: Optional[Union[Mapping[str, str], Iterable[str]]] = None):
  """
  Extracts the implicit windowing information from an element and makes it
  explicit as field(s) in the element itself.

  The following windowing parameter values are supported:

    * `timestamp`: The event timestamp of the current element.
    * `window_start`: The start of the window iff it is an interval window.
    * `window_end`: The (exclusive) end of the window.
    * `window_string`: The string representation of the window.
    * `window_type`: The type of the window as a string.
    * `winodw_object`: The actual window object itself,
        as a Java or Python object.
    * `pane_info`: A schema'd representation of the current pane info, including
        its index, whether it was the last firing, etc.

  As a convenience, a list rather than a mapping of fields may be provided,
  in which case the fields will be named according to the requested values.

  Args:
    fields: A mapping of new field names to various windowing parameters,
      as documented above.  If omitted, defaults to
      `[timestamp, window_start, window_end]`.
  """
  if fields is None:
    fields = ['timestamp', 'window_start', 'window_end']
  if not isinstance(fields, Mapping):
    if isinstance(fields, Iterable) and not isinstance(fields, str):
      fields = {fld: fld for fld in fields}
    else:
      raise TypeError(
          'Fields must be a mapping or iterable of strings, got {fields}')

  existing_fields = named_fields_from_element_type(pcoll.element_type)
  new_fields = []
  for field, value in fields.items():
    if value not in _WINDOWING_INFO_TYPES:
      raise ValueError(
          f'{value} is not a valid windowing parameter; '
          f'must be one of {list(_WINDOWING_INFO_TYPES.keys())}')
    elif field in existing_fields:
      raise ValueError(f'Input schema already has a field named {field}.')
    else:
      new_fields.append((field, _WINDOWING_INFO_TYPES[value]))

  def augment_row(
      row,
      timestamp=beam.DoFn.TimestampParam,
      window=beam.DoFn.WindowParam,
      pane_info=beam.DoFn.PaneInfoParam):
    as_dict = row._asdict()
    for field, value in fields.items():
      as_dict[field] = _WINDOWING_INFO_EXTRACTORS[value](locals())
    return beam.Row(**as_dict)

  return pcoll | beam.Map(augment_row).with_output_types(
      row_type.RowTypeConstraint.from_fields(
          existing_fields + new_fields))  # type: ignore[operator]


def create_mapping_providers():
  # These are MetaInlineProviders because their expansion is in terms of other
  # YamlTransforms, but in a way that needs to be deferred until the input
  # schema is known.
  return [
      yaml_provider.InlineProvider({
          'AssignTimestamps-python': _AssignTimestamps,
          'AssignTimestamps-javascript': _AssignTimestamps,
          'AssignTimestamps-generic': _AssignTimestamps,
          'Explode': _Explode,
          'ExtractWindowingInfo': _ExtractWindowingInfo,
          'Filter-python': _PyJsFilter,
          'Filter-javascript': _PyJsFilter,
          'Filter-generic': _PyJsFilter,
          'MapToFields-python': _PyJsMapToFields,
          'MapToFields-javascript': _PyJsMapToFields,
          'MapToFields-generic': _PyJsMapToFields,
          'Partition-python': _Partition,
          'Partition-javascript': _Partition,
          'Partition-generic': _Partition,
          'StripErrorMetadata': _StripErrorMetadata,
          'ValidateWithSchema': Validate,
      }),
      yaml_provider.SqlBackedProvider({
          'Filter-sql': _SqlFilterTransform,
          'Filter-calcite': _SqlFilterTransform,
          'MapToFields-sql': _SqlMapToFieldsTransform,
          'MapToFields-calcite': _SqlMapToFieldsTransform,
      }),
  ]