Improve error when a user returns an unexpected type from their preprocessing_fn.

Previously for example: "Using a symbolic `tf.Tensor` as a Python `bool` is not allowed in Graph execution".

PiperOrigin-RevId: 466037670

Author: zoyahav

tensorflow_transform/beam/impl.py

@@ -1067,6 +1067,11 @@ def expand(self, dataset):
     # output. This is because if we allowed the output of preprocessing_fn to
     # be empty, we wouldn't be able to determine how many instances to
     # "unbatch" the output into.
+    if not isinstance(structured_outputs, dict):
+      raise ValueError(
+          'A `preprocessing_fn` must return a '
+          'Dict[str, Union[tf.Tensor, tf.SparseTensor, tf.RaggedTensor]]. '
+          f'Got: {structured_outputs}')
     if not structured_outputs:
       raise ValueError('The preprocessing function returned an empty dict')
 

tensorflow_transform/beam/impl_test.py

@@ -4529,6 +4529,18 @@ def preprocessing_fn(inputs):
     self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                           preprocessing_fn, expected_outputs)
 
+  def test_preprocessing_fn_returns_wrong_type(self):
+    with self.assertRaisesRegexp(  # pylint: disable=g-error-prone-assert-raises
+        ValueError, r'A `preprocessing_fn` must return a '
+        r'Dict\[str, Union\[tf.Tensor, tf.SparseTensor, tf.RaggedTensor\]\]. '
+        'Got: Tensor.*'):
+      self.assertAnalyzeAndTransformResults(
+          input_data=[{'f1': 0}],
+          input_metadata=tft.DatasetMetadata.from_feature_spec(
+              {'f1': tf.io.FixedLenFeature([], tf.float32)}),
+          preprocessing_fn=lambda inputs: inputs['f1'],
+          expected_data=None)
+
 
 if __name__ == '__main__':
   tft_unit.main()