(beam) branch master updated: Dask runner: Use bag sqrt scaling (#29366)

[da...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new eedbfe1fce4 Dask runner: Use bag sqrt scaling (#29366)
eedbfe1fce4 is described below

commit eedbfe1fce4ce4bb1a7fad41e53b4c91e5ff6dc0
Author: Jacob Tomlinson <ja...@users.noreply.github.com>
AuthorDate: Tue Dec 5 19:30:02 2023 +0000

    Dask runner: Use bag sqrt scaling (#29366)
    
    * Ensure Dask Runner scales beyond 199 workers
    
    * Add copyright notice
    
    * Fix linting
    
    * Revert copyright change
---
 sdks/python/apache_beam/runners/dask/transform_evaluator.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/sdks/python/apache_beam/runners/dask/transform_evaluator.py b/sdks/python/apache_beam/runners/dask/transform_evaluator.py
index c4aac7f2111..d4d58879b7f 100644
--- a/sdks/python/apache_beam/runners/dask/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/dask/transform_evaluator.py
@@ -24,6 +24,7 @@ TODO(alxr): Translate ops from https://docs.dask.org/en/latest/bag-api.html.
 """
 import abc
 import dataclasses
+import math
 import typing as t
 
 import apache_beam
@@ -59,7 +60,10 @@ class Create(DaskBagOp):
     assert input_bag is None, 'Create expects no input!'
     original_transform = t.cast(_Create, self.transform)
     items = original_transform.values
-    return db.from_sequence(items)
+    return db.from_sequence(
+        items,
+        partition_size=max(
+            1, math.ceil(math.sqrt(len(items)) / math.sqrt(100))))
 
 
 class ParDo(DaskBagOp):