Apache Flink: Map function over values of a KeyedStream, without losing keyed state

I have a keyed stream:

KeyedStream<String, Data> keyed = 
env.addSource(...)
   .map(new Parser())
   .keyBy(i -> i.key)

I want to run a state-full process function on my stream; but the process will return a normal un-keyed stream that cause losing KeyedStream and force my to call keyBy again:

SingleOutputStreamOperator<Data> unkeyed = keyed.process(new Function)
KeyedStream<String, Data> keyedAgain = keyed.keyBy(i -> i.key)

• Is there any way to call a map function over values on a KeyedStream?

• How I can run a state-full function without using a low-level API like ProcessFunction that cause unwanted cost: a new sub-task at runtime, plus re-shuffeling my data?