[Bug]: ParDo._get_key_and_window_coder fails to filter side inputs when determining main input for stateful DoFns #40374

main_input = list(set(named_inputs.keys()) - set(self.side_inputs))[0]

import apache_beam as beam
from apache_beam.transforms.userstate import ReadModifyWriteStateSpec, TimerSpec, on_timer
from apache_beam.utils.timestamp import Timestamp
from apache_beam.portability.api import beam_runner_api_pb2

STATE = ReadModifyWriteStateSpec('state', beam.coders.VarIntCoder())
TIMER = TimerSpec('timer', beam.TimeDomain.WATERMARK)

class StatefulFn(beam.DoFn):
    def process(
        self,
        element: tuple[str, int],
        state=beam.DoFn.StateParam(STATE),
        timer=beam.DoFn.TimerParam(TIMER),
        side: dict = None,
    ):
        timer.set(Timestamp.now() + 60)
        yield element

    @on_timer(TIMER)
    def on_timer(self):
        yield 1

p = beam.Pipeline()
main_in = p | "Main" >> beam.Create([("key1", 1)])
side_in = p | "Side" >> beam.Create([("s", 1)])
side_dict = beam.pvalue.AsDict(side_in)

main_in | beam.ParDo(StatefulFn(), side=side_dict)
proto = p.to_runner_api()

# Depending on PYTHONHASHSEED, timer family specs will intermittently use FastPrimitivesCoder for key_coder instead of StrUtf8Coder / KVCoder

def _get_key_and_window_coder(self, named_inputs):
    if named_inputs is None or not self._signature.is_stateful_dofn():
        return None, None

    side_input_pcolls = {si.pvalue for si in (self.side_inputs or ())}
    main_input_candidates = [
        k
        for k, pcoll in named_inputs.items()
        if pcoll not in side_input_pcolls
        and not str(k).startswith(core.SIDE_INPUT_PREFIX)
    ]
    if main_input_candidates:
        main_input = main_input_candidates[0]
    else:
        main_input = list(named_inputs.keys())[0]

    input_pcoll = named_inputs[main_input]
    kv_type_hint = input_pcoll.element_type
    if kv_type_hint and kv_type_hint != typehints.Any:
        coder = coders.registry.get_coder(kv_type_hint)
        if not coder.is_kv_coder():
            raise ValueError(
                f"Input elements to the transform {self} with stateful DoFn must be "
                "key-value pairs."
            )
        key_coder = coder.key_coder()
    else:
        key_coder = coders.registry.get_coder(typehints.Any)
    window_coder = input_pcoll.windowing.windowfn.get_window_coder()
    return key_coder, window_coder