Refine length bucketing docs and fix boundary inclusivity

Expands parameter documentation for clarity and replaces bisect_left with bisect_right to ensure bucket boundaries are inclusive on the lower bound. Updates util_test.py assertions accordingly.

Author: Eliaaazzz

sdks/python/apache_beam/ml/inference/base.py

@@ -178,8 +178,8 @@ def __init__(
       max_batch_duration_secs: Optional[int] = None,
       max_batch_weight: Optional[int] = None,
       element_size_fn: Optional[Callable[[Any], int]] = None,
-      length_fn: Optional[Callable[[Any], int]] = None,
-      bucket_boundaries: Optional[list[int]] = None,
+      batch_length_fn: Optional[Callable[[Any], int]] = None,
+      batch_bucket_boundaries: Optional[list[int]] = None,
       large_model: bool = False,
       model_copies: Optional[int] = None,
       **kwargs):
@@ -192,11 +192,17 @@ def __init__(
         before emitting; used in streaming contexts.
       max_batch_weight: the maximum weight of a batch. Requires element_size_fn.
       element_size_fn: a function that returns the size (weight) of an element.
-      length_fn: a callable mapping an element to its length. When set with
-        max_batch_duration_secs, enables length-aware bucketed keying so
-        elements of similar length are batched together.
-      bucket_boundaries: sorted list of positive boundary values for length
-        bucketing. Requires length_fn.
+      batch_length_fn: a callable mapping an element to its length (int). When
+        set together with max_batch_duration_secs, enables length-aware bucketed
+        keying so that elements of similar length are batched together, reducing
+        padding waste for variable-length inputs. Bucket assignment uses
+        bisect_right so boundaries are lower-inclusive: e.g., for boundaries
+        [10, 50], buckets are (-inf, 10), [10, 50), [50, inf).
+      batch_bucket_boundaries: a sorted list of positive boundary values for
+        length bucketing. Boundaries are lower-inclusive (bisect_right
+        semantics): bucket i covers lengths in [boundaries[i-1], boundaries[i]).
+        Requires batch_length_fn. Defaults to [16, 32, 64, 128, 256, 512] when
+        batch_length_fn is set.
       large_model: set to true if your model is large enough to run into
         memory pressure if you load multiple copies.
       model_copies: The exact number of models that you would like loaded
@@ -216,10 +222,10 @@ def __init__(
       self._batching_kwargs['max_batch_weight'] = max_batch_weight
     if element_size_fn is not None:
       self._batching_kwargs['element_size_fn'] = element_size_fn
-    if length_fn is not None:
-      self._batching_kwargs['length_fn'] = length_fn
-    if bucket_boundaries is not None:
-      self._batching_kwargs['bucket_boundaries'] = bucket_boundaries
+    if batch_length_fn is not None:
+      self._batching_kwargs['length_fn'] = batch_length_fn
+    if batch_bucket_boundaries is not None:
+      self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries
     self._large_model = large_model
     self._model_copies = model_copies
     self._share_across_processes = large_model or (model_copies is not None)

sdks/python/apache_beam/ml/inference/base_test.py

@@ -2279,38 +2279,40 @@ def test_max_batch_duration_secs_only(self):
 
     self.assertEqual(kwargs, {'max_batch_duration_secs': 60})
 
-  def test_length_fn_and_bucket_boundaries(self):
-    """length_fn and bucket_boundaries are passed through to kwargs."""
+  def test_batch_length_fn_and_batch_bucket_boundaries(self):
+    """batch_length_fn and batch_bucket_boundaries passed through to kwargs."""
     handler = FakeModelHandlerForBatching(
-        length_fn=len, bucket_boundaries=[16, 32, 64])
+        batch_length_fn=len, batch_bucket_boundaries=[16, 32, 64])
     kwargs = handler.batch_elements_kwargs()
 
     self.assertIs(kwargs['length_fn'], len)
     self.assertEqual(kwargs['bucket_boundaries'], [16, 32, 64])
 
-  def test_length_fn_only(self):
-    """length_fn alone is passed through without bucket_boundaries."""
-    handler = FakeModelHandlerForBatching(length_fn=len)
+  def test_batch_length_fn_only(self):
+    """batch_length_fn alone is passed through without bucket_boundaries."""
+    handler = FakeModelHandlerForBatching(batch_length_fn=len)
     kwargs = handler.batch_elements_kwargs()
 
     self.assertIs(kwargs['length_fn'], len)
     self.assertNotIn('bucket_boundaries', kwargs)
 
-  def test_bucket_boundaries_without_length_fn(self):
-    """Passing bucket_boundaries without length_fn should fail in BatchElements.
+  def test_batch_bucket_boundaries_without_batch_length_fn(self):
+    """Passing batch_bucket_boundaries without batch_length_fn should fail in
+    BatchElements.
 
     Note: ModelHandler.__init__ doesn't validate this; the error is raised
     by BatchElements when batch_elements_kwargs are used."""
-    handler = FakeModelHandlerForBatching(bucket_boundaries=[10, 20])
+    handler = FakeModelHandlerForBatching(batch_bucket_boundaries=[10, 20])
     kwargs = handler.batch_elements_kwargs()
     # The kwargs are stored, but BatchElements will reject them
     self.assertEqual(kwargs['bucket_boundaries'], [10, 20])
     self.assertNotIn('length_fn', kwargs)
 
   def test_batching_kwargs_none_values_omitted(self):
-    """None values for length_fn and bucket_boundaries are not in kwargs."""
+    """None values for batch_length_fn and batch_bucket_boundaries are not in
+    kwargs."""
     handler = FakeModelHandlerForBatching(
-        min_batch_size=5, length_fn=None, bucket_boundaries=None)
+        min_batch_size=5, batch_length_fn=None, batch_bucket_boundaries=None)
     kwargs = handler.batch_elements_kwargs()
     self.assertNotIn('length_fn', kwargs)
     self.assertNotIn('bucket_boundaries', kwargs)

sdks/python/apache_beam/transforms/util.py

@@ -1223,7 +1223,9 @@ def setup(self):
         load_shared_key, "WithLengthBucketKey").key
 
   def _get_bucket(self, length):
-    return bisect.bisect_left(self._bucket_boundaries, length)
+    # bisect_right: boundaries are lower-inclusive.
+    # e.g., for boundaries [10, 50], buckets are (-inf, 10), [10, 50), [50, inf)
+    return bisect.bisect_right(self._bucket_boundaries, length)
 
   def process(self, element):
     length = self._length_fn(element)
@@ -1291,14 +1293,14 @@ class BatchElements(PTransform):
     record_metrics: (optional) whether or not to record beam metrics on
         distributions of the batch size. Defaults to True.
     length_fn: (optional) a callable mapping an element to its length (int).
-        When set together with max_batch_duration_secs, enables length-aware
-        bucketed keying on the stateful path so that elements of similar length
-        are routed to the same batch, reducing padding waste.
+        When set together with bucket_boundaries, enables length-aware bucketed
+        keying on the stateful path so that elements of similar length are
+        routed to the same batch, reducing padding waste.
     bucket_boundaries: (optional) a sorted list of positive boundary values
-        for length bucketing. Elements with length < boundaries[i] go to
-        bucket i; overflow goes to bucket len(boundaries). Defaults to
-        [16, 32, 64, 128, 256, 512] when length_fn is set. Requires
-        length_fn.
+        for length bucketing. Boundaries are lower-inclusive (bisect_right
+        semantics): e.g., for boundaries [10, 50], buckets are (-inf, 10),
+        [10, 50), [50, inf). Defaults to [16, 32, 64, 128, 256, 512] when
+        length_fn is set. Requires length_fn.
   """
   _DEFAULT_BUCKET_BOUNDARIES = [16, 32, 64, 128, 256, 512]
 

sdks/python/apache_beam/transforms/util_test.py

@@ -1030,13 +1030,15 @@ def test_length_bucket_assignment(self):
     """WithLengthBucketKey assigns correct bucket indices."""
     boundaries = [10, 50, 100]
     dofn = util.WithLengthBucketKey(length_fn=len, bucket_boundaries=boundaries)
-    # bisect_left: length < 10 -> bucket 0, 10 <= length < 50 -> bucket 1, etc.
+    # bisect_right: boundaries are lower-inclusive.
+    # e.g., for boundaries [10, 50, 100], buckets are:
+    #   (-inf, 10), [10, 50), [50, 100), [100, inf)
     self.assertEqual(dofn._get_bucket(5), 0)
-    self.assertEqual(dofn._get_bucket(10), 0)
+    self.assertEqual(dofn._get_bucket(10), 1)
     self.assertEqual(dofn._get_bucket(11), 1)
-    self.assertEqual(dofn._get_bucket(50), 1)
+    self.assertEqual(dofn._get_bucket(50), 2)
     self.assertEqual(dofn._get_bucket(51), 2)
-    self.assertEqual(dofn._get_bucket(100), 2)
+    self.assertEqual(dofn._get_bucket(100), 3)
     self.assertEqual(dofn._get_bucket(101), 3)
     self.assertEqual(dofn._get_bucket(999), 3)