[Bug]: Type inference failing for Python SDK with External transforms and beam.Row

[rahuliyer95]

What happened?

When using expansion service with Beam SDK 2.41.0, the type inference is failing with the following exception,

Traceback (most recent call last):
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/__main__.py", line 87, in <module>
    bootstrap_pex(__entry_point__)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex_bootstrapper.py", line 591, in bootstrap_pex
    pex.PEX(entry_point).execute()
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex.py", line 528, in execute
    sys.exit(self._wrap_coverage(self._wrap_profiling, self._execute))
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex.py", line 435, in _wrap_coverage
    return runner(*args)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex.py", line 466, in _wrap_profiling
    return runner(*args)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex.py", line 556, in _execute
    EntryPoint.parse("run = {}".format(self._pex_info.entry_point))
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex.py", line 739, in execute_entry
    return self.execute_module(entry_point.module)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.bootstrap/pex/pex.py", line 747, in execute_module
    runpy.run_module(module_name, run_name="__main__", alter_sys=True)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 205, in run_module
    return _run_module_code(code, init_globals, run_name, mod_spec)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 96, in _run_module_code
    mod_name, mod_spec, pkg_name, script_name)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/twitter/udp/integration_tests/jobs/dal_read_job.py", line 180, in <module>
    main()
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/twitter/udp/integration_tests/jobs/dal_read_job.py", line 157, in main
    >> dal_rows.ReadRows("tweetsource-public", "public_tweets", opts.date).with_output_types(beam.Row)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/ptransform.py", line 1095, in __ror__
    return self.transform.__ror__(pvalueish, self.label)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/ptransform.py", line 617, in __ror__
    result = p.apply(self, pvalueish, label)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/pipeline.py", line 663, in apply
    return self.apply(transform, pvalueish)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/pipeline.py", line 709, in apply
    pvalueish_result = self.runner.apply(transform, pvalueish, self._options)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/dataflow/dataflow_runner.py", line 141, in apply
    return super().apply(transform, input, options)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/runner.py", line 185, in apply
    return m(transform, input, options)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/runner.py", line 215, in apply_PTransform
    return transform.expand(input)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/external.py", line 547, in expand
    pcoll_id in self._expanded_transform.outputs.items()
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/external.py", line 547, in <dictcomp>
    pcoll_id in self._expanded_transform.outputs.items()
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/pipeline_context.py", line 115, in get_by_id
    self._id_to_proto[id], self._pipeline_context)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/pvalue.py", line 210, in from_runner_api
    element_type=context.element_type_from_coder_id(proto.coder_id),
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/pipeline_context.py", line 269, in element_type_from_coder_id
    self.coders[coder_id].to_type_hint())
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/pipeline_context.py", line 163, in __getitem__
    return self.get_by_id(id)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/pipeline_context.py", line 115, in get_by_id
    self._id_to_proto[id], self._pipeline_context)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/coders/coders.py", line 385, in from_runner_api
    context)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/coders/row_coder.py", line 111, in from_runner_api_parameter
    return RowCoder(schema)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/coders/row_coder.py", line 63, in __init__
    self._type_hint = named_tuple_from_schema(self.schema)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/typehints/schemas.py", line 462, in named_tuple_from_schema
    schema_registry=schema_registry)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/typehints/schemas.py", line 179, in typing_from_runner_api
    schema_registry=schema_registry).typing_from_runner_api(fieldtype_proto)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/typehints/schemas.py", line 407, in typing_from_runner_api
    field_py_type = self.typing_from_runner_api(field.type)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmps6_sxbu8/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/typehints/schemas.py", line 365, in typing_from_runner_api
    return Optional[base]
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/typing.py", line 254, in inner
    return func(*args, **kwds)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/typing.py", line 357, in __getitem__
    arg = _type_check(parameters, "Optional[t] requires a single type.")
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/typing.py", line 142, in _type_check
    raise TypeError(f"{msg} Got {arg!r:.100}.")
TypeError: Optional[t] requires a single type. Got Row(location=typing.Union[typing.Sequence[str], NoneType], empty__=typing.Union[bool, NoneType]).

This is how we are using the external transform,

import apache_beam as beam
from apache_beam.options.pipeline_options import GoogleCloudOptions, PipelineOptions
from apache_beam.transforms.external import ExpansionAndArtifactRetrievalStub
import grpc
import typing

External = typing.NamedTuple(
  "External",
  [
    ("id", str),
  ],
)

creds = grpc.ssl_channel_credentials()
channel = grpc.secure_channel("localhost:10001", creds)
grpc.channel_ready_future(channel).result()
stub = ExpansionAndArtifactRetrievalStub(channel)
with beam.Pipeline(options=PipelineOptions()) as p:
  p | "Read" >> beam.ExternalTransform("URN", NamedTupleBasedPayloadBuilder(External("123")), stub)

The same code works fine with SDK 2.40.0

Issue Priority

Priority: 1

Issue Component

Component: sdk-py-core

[rahuliyer95]

Trying to run with --no_pipeline_type_check fails with the following exception,

Traceback (most recent call last):
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/__main__.py", line 87, in <module>
    bootstrap_pex(__entry_point__)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex_bootstrapper.py", line 591, in bootstrap_pex
    pex.PEX(entry_point).execute()
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex.py", line 528, in execute
    sys.exit(self._wrap_coverage(self._wrap_profiling, self._execute))
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex.py", line 435, in _wrap_coverage
    return runner(*args)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex.py", line 466, in _wrap_profiling
    return runner(*args)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex.py", line 556, in _execute
    EntryPoint.parse("run = {}".format(self._pex_info.entry_point))
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex.py", line 739, in execute_entry
    return self.execute_module(entry_point.module)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.bootstrap/pex/pex.py", line 747, in execute_module
    runpy.run_module(module_name, run_name="__main__", alter_sys=True)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 205, in run_module
    return _run_module_code(code, init_globals, run_name, mod_spec)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 96, in _run_module_code
    mod_name, mod_spec, pkg_name, script_name)
  File "/opt/twitter_mde/package/eepython37/886055eae06dbd005860a8b21c6f43355dd92462660a6eda8e580552513a78a2/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/twitter/udp/integration_tests/jobs/dal_read_job.py", line 180, in <module>
    main()
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/twitter/udp/integration_tests/jobs/dal_read_job.py", line 162, in main
    | "PublicTweets.Rows.WriteText"
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/ptransform.py", line 1095, in __ror__
    return self.transform.__ror__(pvalueish, self.label)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/ptransform.py", line 617, in __ror__
    result = p.apply(self, pvalueish, label)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/pipeline.py", line 663, in apply
    return self.apply(transform, pvalueish)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/pipeline.py", line 709, in apply
    pvalueish_result = self.runner.apply(transform, pvalueish, self._options)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/dataflow/dataflow_runner.py", line 141, in apply
    return super().apply(transform, input, options)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/runner.py", line 185, in apply
    return m(transform, input, options)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/runners/runner.py", line 215, in apply_PTransform
    return transform.expand(input)
  File "/private/var/folders/yt/m4jhm2d5197c62c7v0kx77fc0000gp/T/tmp9busrhl0/.deps/apache_beam-2.41.0-py3-none-any.whl/apache_beam/transforms/external.py", line 506, in expand
    if self._type_hints.output_types:
AttributeError: 'ExternalTransform' object has no attribute '_type_hints'

[chamikaramj]

cc: @TheNeuralBit

[TheNeuralBit]

Thanks for the ping, this definitely looks bad. I'm not sure how existing tests didn't catch it. I would expect sql_test.py to trigger it at least.

[TheNeuralBit]

Ah it looks like this issue may be specifically for an optional (nullable), nested row field, which could very well be a gap in our testing. Like:

Row {
  foo: INT
  bar: OPTIONAL[ROW { baz: STR }]
}

Do you know the schema returned by the external transform in this case @rahuliyer95?

Regardless, this is almost certainly caused by #22679, we could roll that back but we'd need to do a point release to get it in 2.41.

[chamikaramj]

Thanks Brian. If this is for a very specific case, I think we can just document that as a known issues for 2.41.0 release and fix for the 2.42.0 release (cut in two weeks). Also we should add the missing test.

[TheNeuralBit]

Hm actually the breakage is due to #22066, which is tougher to rollback.

The core issue here is that PR started to mix beam typehints (RowTypeConstraint) and native typehints (typing.Optional). In all our test cases this was OK because the RowTypeConstraint was either the top-level, or if nested it was not nullable/optional, so it didn't need to be wrapped in a typing.Optional.

[rahuliyer95]

@TheNeuralBit I attached a debug point on line 362 in schemas.py and here is the value for fieldtype_proto

nullable: true
row_type {
  schema {
    fields {
      name: "location"
      type {
        nullable: true
        array_type {
          element_type {
            atomic_type: STRING
          }
        }
      }
    }
    fields {
      name: "empty__"
      type {
        nullable: true
        atomic_type: BOOLEAN
      }
      id: 1
      encoding_position: 1
    }
    id: "be01b00e-d749-471b-bb17-4b4358d45070"
  }
}