Add validation to check for non-ints from partitionfns

sdks/python/apache_beam/transforms/core.py

@@ -3573,6 +3573,10 @@ class ApplyPartitionFnFn(DoFn):
     """A DoFn that applies a PartitionFn."""
     def process(self, element, partitionfn, n, *args, **kwargs):
       partition = partitionfn.partition_for(element, n, *args, **kwargs)
+      if isinstance(partition, bool) or not isinstance(partition, int):
+        raise ValueError(
+            f"PartitionFn yielded a '{type(partition).__name__}' "
+            "when it should only yield integers")
       if not 0 <= partition < n:
         raise ValueError(
             'PartitionFn specified out-of-bounds partition index: '

sdks/python/apache_beam/transforms/core_test.py

@@ -162,6 +162,20 @@ def test_dofn_with_implicit_return_none_return_without_value(self):
 
 
 class PartitionTest(unittest.TestCase):
+  def test_partition_with_bools(self):
+    with pytest.raises(
+        (ValueError, RuntimeError),
+        match=
+        r"PartitionFn yielded a '([^']*)' when it should only yield integers"):
+      # Check for RuntimeError too since the portable runner casts
+      # all exceptions to RuntimeError
+      invalid_inputs = [True, 1.2, 'string', None]
+      for input_value in invalid_inputs:
+        with beam.testing.test_pipeline.TestPipeline() as p:
+          _ = (
+              p | beam.Create([input_value])
+              | beam.Partition(lambda x, _: x, 2))
+
   def test_partition_boundedness(self):
     def partition_fn(val, num_partitions):
       return val % num_partitions