[tf.data] Internal change to nested `Dataset` handling in functions. No functional change.

This change adds experimental support in `StructuredFunctionWrapper` for taking datasets as function arguments and returning them, which is a stepping stone to creating datasets of datasets. However, the public API for such datasets (in particular how introspection using `Dataset.output_classes` etc. would work) is not settled, so this change only uses the new support in the implementation of `Dataset.flat_map()` and `tf.contrib.data.group_by_window()`.

The change also folds some code into its sole point of use in `StructuredFunctionWrapper`.

PiperOrigin-RevId: 201613301