apache/beam · error · KeyError

Transform is not registered with a supported type. Please…

Error message

Transform {class_name} is not registered with a supported type. Please register the transform with a supported type using register_input_dtype decorator.

What it means

MLTransform maps columns to input dtypes via the _EXPECTED_TYPES registry, populated by the @register_input_dtype decorator. If a transform's class was never registered with a supported input dtype, the column-type mapping cannot proceed and KeyError is raised.

Solutions

Decorate your transform class with @register_input_dtype(dtype) so it lands in _EXPECTED_TYPES.
Use a built-in supported transform instead of the custom class.
If a class was renamed, register the new class name or revert the rename.

Example fix

// before
class MyEmbeddings(EmbeddingsHandler):
    ...
// after
@register_input_dtype(str)
class MyEmbeddings(EmbeddingsHandler):
    ...

Defensive patterns

Strategy: try-catch

Validate before calling

from apache_beam.ml.transforms.handlers import _EXPECTED_TYPES
assert type(my_transform).__name__ in _EXPECTED_TYPES

Try / catch

try:
    result = (pcoll | MLTransform(...))
except KeyError as e:
    if 'not registered with a supported' in str(e):
        register_input_dtype(str)(MyTransform)
        result = (pcoll | MLTransform(...))

Prevention

Apply @register_input_dtype to every custom transform
Check _EXPECTED_TYPES before composing new transforms

When it happens

Trigger: Using a custom (or third-party) transformation class inside MLTransform without decorating it with @register_input_dtype(input_type=...); class renamed so the registry lookup by class name misses.

Common situations: Writing a custom RunInference-style or custom embedding handler and plugging it into MLTransform; upgrading Beam where a transform's class name changed.

Related errors

AI-assisted analysis of apache/beam@12126d8942 (2026-09-13). Data as JSON: /api/errors/177a1e2e1cee8c4e. Report an issue: GitHub.

Appendix: source

Thrown at sdks/python/apache_beam/ml/transforms/handlers.py:210

      # sometimes a numpy type can be provided as np.dtype('int64').
      # convert numpy.dtype to numpy type since both are same.
      for name, typ in inferred_types.items():
        if isinstance(typ, np.dtype):
          inferred_types[name] = typ.type

      return inferred_types
    except:  # pylint: disable=bare-except
      return {}

  def _map_column_names_to_types_from_transforms(self):
    column_type_mapping = {}
    for transform in self.transforms:
      for col in transform.columns:
        if col not in column_type_mapping:
          # we just need to dtype of first occurance of column in transforms.
          class_name = transform.__class__.__name__
          if class_name not in _EXPECTED_TYPES:
            raise KeyError(
                f"Transform {class_name} is not registered with a supported "
                "type. Please register the transform with a supported type "
                "using register_input_dtype decorator.")
          column_type_mapping[col] = _EXPECTED_TYPES[
              transform.__class__.__name__]
    return column_type_mapping

  def get_raw_data_feature_spec(
      self, input_types: dict[str, type]) -> dict[str, tf.io.VarLenFeature]:
    """
    Return a DatasetMetadata object to be used with
    tft_beam.AnalyzeAndTransformDataset.
    Args:
      input_types: A dictionary of column names and types.
    Returns:
      A DatasetMetadata object.
    """
    raw_data_feature_spec = {}

View on GitHub (pinned to 12126d8942)