apache/beam · error · NotImplementedError

https://github.com/apache/beam/issues/21653: Per-key…

Error message

https://github.com/apache/beam/issues/21653: Per-key process_batch

What it means

Beam's batched SDF processing (process_batch) does not yet support a per-key KeyParam injection, so a DoFn whose process_batch declares KeyParam raises this NotImplementedError referencing beam issue 21653. This is an intentional, unimplemented-feature guard, not a data error.

Solutions

Remove KeyParam from process_batch and derive the key inside the method from the elements (they are KVs)
Fall back to per-element process() until BEAM-21653 is implemented
Pre-group/key elements explicitly and use process() with StateParam if per-key access is essential

Example fix

# before
def process_batch(self, batch, key=DoFn.KeyParam):
    ...
# after
def process_batch(self, batch):
    for k, v in batch.values:
        ...  # extract key from each KV element

Defensive patterns

Strategy: fallback

Validate before calling

import inspect

def batch_supports_key(dofn):
    src = inspect.getsource(dofn.process_batch)
    return 'DoFn.KeyParam' not in src

Try / catch

try:
    invoker.invoke_process_batch(batch)
except NotImplementedError as e:
    if '21653' in str(e):
        logging.warning('process_batch lacks KeyParam support; falling back to process()')
        for wb in batch.split():
            invoker.invoke_process(wb)

Prevention

Do not use KeyParam in process_batch until BEAM-21653 lands
Extract keys from KV elements inside process_batch instead
Track the Beam issue for feature availability before migrating

When it happens

Trigger: Declaring core.DoFn.KeyParam as a parameter of a process_batch method and invoking it via DoFnInvoker.invoke_process_batch on windowed batches.

Common situations: Porting an existing per-element process() (with KeyParam) to process_batch for performance; using batched stateful DoFns on runners that route through the batched invoker.

Related errors

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

Appendix: source

Thrown at sdks/python/apache_beam/runners/common.py:1121

      side_inputs = [si[window] for si in self.side_inputs]
      side_inputs.extend(additional_args)
      args_for_process_batch, kwargs_for_process_batch = (
          util.insert_values_in_args(
              self.args_for_process_batch,
              self.kwargs_for_process_batch,
              side_inputs,
          )
      )
      if not self.recalculate_window_args:
        self.args_for_process_batch, self.kwargs_for_process_batch = (
            args_for_process_batch, kwargs_for_process_batch)
        self.has_cached_window_batch_args = True

    for i, p in self.placeholders_for_process_batch:
      if core.DoFn.ElementParam == p:
        args_for_process_batch[i] = windowed_batch.values
      elif core.DoFn.KeyParam == p:
        raise NotImplementedError(
            'https://github.com/apache/beam/issues/21653: Per-key process_batch'
        )
      elif core.DoFn.WindowParam == p:
        args_for_process_batch[i] = window
      elif core.DoFn.TimestampParam == p:
        args_for_process_batch[i] = windowed_batch.timestamp
      elif core.DoFn.PaneInfoParam == p:
        assert isinstance(windowed_batch, HomogeneousWindowedBatch)
        args_for_process_batch[i] = windowed_batch.pane_info
      elif isinstance(p, core.DoFn.StateParam):
        raise NotImplementedError(
            "https://github.com/apache/beam/issues/21653: "
            "Per-key process_batch")
      elif isinstance(p, core.DoFn.TimerParam):
        raise NotImplementedError(
            "https://github.com/apache/beam/issues/21653: "
            "Per-key process_batch")
      elif isinstance(p, core.DoFn.BundleContextParam):

View on GitHub (pinned to 12126d8942)