Add batch support to generator framework

boundary_layer/builders/templates/generator_operator.j2

@@ -29,23 +29,61 @@ def {{ item_name_builder }}(index, item):
  latter would discard any default task args, expecting them to be filled-in
  by airflow, while in fact airflow would not fill them in at all. #}
 {% set properties = node.resolved_properties.values %}
-for (index, item) in enumerate({{ iterable_builder }}(
+
+{% set all_items = (node.name + '_all_items') | sanitize_operator_name %}
+{{ all_items }} = {{ iterable_builder }}(
 {% for arg in builder_args %}
 {% if arg in properties %}
             {{ arg }} = {{ properties[arg] | format_value }},
 {% endif %}
 {% endfor %}
-        )):
+        )
+
+{% if node.batching_enabled %}
+{# Generate code for batched situations #}
+{% set batch_name_builder = (node.name + '_batch_name_builder') | sanitize_operator_name %}
+def {{ batch_name_builder }}(index, items):
+    return 'batch_%d_%d' % (index, len(items))
+
+{# TODO: Import this from some util module when such functionality is possible #}
+def generator_helper_filter_with_blocklist(items, item_name_builder, blocklist):
+    def not_in_blocklist(index, item):
+        item_name = item_name_builder(index, item)
+        return not any(re.match(i, item_name) for i in blocklist)
+
+    filtered = filter(lambda (index, item): not_in_blocklist(index, item), enumerate(items))
+
+    return map(lambda t: t[1], filtered)
 
+{# TODO: Import this from some util module when such functionality is possible #}
+{# Borrowed from: https://stackoverflow.com/a/312464 #}
+def generator_helper_grouped_list(l, n):
+    for i in range(0, len(l), n):
+        yield l[i:i + n]
+
+{% set filtered = (node.name + '_filtered') | sanitize_operator_name %}
+{{ filtered }} = generator_helper_filter_with_blocklist({{ all_items }}, {{ item_name_builder }}, {{ blocklist }})
+
+for (index, items) in enumerate(generator_helper_grouped_list({{ filtered }}, {{ node.batching.batch_size }})):
+    batch_name = {{ batch_name_builder }}(index, items)
+
+{% set item_input = 'items' %}
+{% set name_input = 'batch_name' %}
+{% else %}
+{# Generate code for non-batched situations #}
+for (index, item) in enumerate({{ all_items }}):
     item_name = {{ item_name_builder }}(index, item)
     blocklist_match = any(re.match(i, item_name) for i in {{ blocklist }})
     if blocklist_match:
         continue
 
+{% set item_input = 'item' %}
+{% set name_input = 'item_name' %}
+{% endif %}
     {{ node.target | sanitize_operator_name }}_builder(
         index = index,
-        item = item,
-        item_name = item_name,
+        {{ item_input }} = {{ item_input }},
+        {{ name_input }} = {{ name_input }},
         dag = dag,
         upstream_dependencies = {{ upstream_dependencies | sanitize_operator_name | verbatim | format_value }},
         downstream_dependencies = {{ downstream_dependencies | sanitize_operator_name | verbatim | format_value }})

boundary_layer/builders/templates/generator_preamble.j2

@@ -13,10 +13,17 @@ You may obtain a copy of the License at
     See the License for the specific language governing permissions and
     limitations under the License.
 #}
+{% if referring_node.batching_enabled %}
+    {%- set item_input = 'items' %}
+    {%- set name_input = 'batch_name' %}
+{% else %}
+    {%- set item_input = 'item' %}
+    {%- set name_input = 'item_name' %}
+{% endif %}
 def {{ generator_operator_name | sanitize_operator_name }}_builder(
         index,
-        item,
-        item_name,
+        {{ item_input }},
+        {{ name_input }},
         dag,
         upstream_dependencies,
         downstream_dependencies):

boundary_layer/registry/types/generator.py

@@ -28,6 +28,14 @@ class GeneratorNode(SubdagNode):
     def regex_blocklist(self):
         return self.item.get('regex_blocklist', ())
 
+    @property
+    def batching(self):
+        return self.item.get('batching', {})
+
+    @property
+    def batching_enabled(self):
+        return not self.batching.get('disabled') if self.batching else False
+
 
 class GeneratorRegistry(ConfigFileRegistry):
     node_cls = GeneratorNode

boundary_layer/registry/types/operator.py

@@ -422,8 +422,18 @@ def _build_task_id(self, execution_context):
             return base_name
 
         suffix_mode = execution_context.referrer.item.get('auto_task_id_mode')
-        if not suffix_mode or suffix_mode == 'item_name':
-            return base_name + '-<<item_name>>'
+        batching_enabled = execution_context.referrer.batching_enabled
+        # Validate suffix_mode based on batching config
+        if batching_enabled and suffix_mode == 'item_name':
+            raise Exception(
+                'Cannot use `item_name` for auto_task_id_mode when batching is enabled')
+        elif not batching_enabled and suffix_mode == 'batch_name':
+            raise Exception(
+                'Cannot use `batch_name` for auto_task_id_mode when batching is disabled')
+
+        name_var = 'batch_name' if batching_enabled else 'item_name'
+        if not suffix_mode or suffix_mode == name_var:
+            return base_name + '-<<' + name_var + '>>'
         elif suffix_mode == 'index':
             return base_name + '-<<str(index)>>'
 

boundary_layer/schemas/dag.py

@@ -35,9 +35,15 @@ class ReferenceSchema(OperatorSchema):
     target = fields.String(required=True)
 
 
+class BatchingSchema(StrictSchema):
+    disabled = fields.Boolean()
+    batch_size = fields.Integer(required=True)
+
+
 class GeneratorSchema(ReferenceSchema):
     auto_task_id_mode = fields.String()
     regex_blocklist = fields.List(fields.String())
+    batching = fields.Nested(BatchingSchema)
 
     @validates_schema
     def check_task_id_mode(self, data):

test/test_generators.py

@@ -1,6 +1,13 @@
-from boundary_layer.registry import NodeTypes
-from boundary_layer.containers import ExecutionContext
+import copy
+import re
+import yaml
+
 from boundary_layer import plugins
+from boundary_layer.builders import PrimaryDagBuilder
+from boundary_layer.containers import ExecutionContext
+from boundary_layer.registry import NodeTypes
+from boundary_layer.registry.types.generator import GeneratorNode
+from boundary_layer.schemas.internal.generators import GeneratorSpecSchema
 
 
 def test_default_param_filler():
@@ -25,3 +32,186 @@ def test_default_param_filler():
         'timeout_sec': 5,
         'headers': {}
     }
+
+
+# Tests for batching functionality
+
+BASE_GENERATOR_CONFIG = {
+    'name': 'test_generator',
+    'type': 'list_generator',
+    'target': 'some_target',
+    'properties': {
+        'items': ['a', 'b', 'c', 'd', 'e', 'f', 'g']
+    }
+}
+
+GENERATOR_CONFIG_YAML = """
+name: list_generator
+iterator_builder_method_code: return items
+item_name_builder_code: return item
+parameters_jsonschema:
+    properties:
+        items:
+            type: array
+            items:
+                type: string
+    additionalProperties: false
+    required:
+        - items
+"""
+
+
+class BatchingTestHelper(object):
+    """
+    Helper class to reduce code required to test code generation under different cases.
+    """
+
+    def __init__(self):
+        self.builder = PrimaryDagBuilder(None, None, None, None)
+        self.generator_spec_schema = GeneratorSpecSchema().load(yaml.load(GENERATOR_CONFIG_YAML))
+
+    def build_generator_node(self, batching_config):
+        node_config = copy.deepcopy(BASE_GENERATOR_CONFIG)
+        if batching_config is not None:
+            node_config['batching'] = batching_config
+
+        return GeneratorNode(config=self.generator_spec_schema.data, item=node_config)
+
+    def run_preamble_template_test(self, batching_config):
+        node = self.build_generator_node(batching_config)
+        template = self.builder.get_jinja_template('generator_preamble.j2')
+
+        rendered = template.render(
+            generator_operator_name='foo',
+            referring_node=node
+        )
+
+        items_batch_name_regex = re.compile(r'\s+items,\s+batch_name,')
+        item_item_name_regex = re.compile(r'\s+item,\s+item_name,')
+
+        items_batch_name_match = items_batch_name_regex.search(rendered)
+        item_item_name_match = item_item_name_regex.search(rendered)
+
+        return {
+            'items_batch_name': items_batch_name_match,
+            'item_item_name': item_item_name_match
+        }
+
+    def run_operator_template_test(self, batching_config):
+        node = self.build_generator_node(batching_config)
+        template = self.builder.get_jinja_template('generator_operator.j2')
+
+        node.resolve_properties(
+            execution_context=ExecutionContext(None, {}),
+            default_task_args={},
+            base_operator_loader=None,
+            preprocessor_loader=None
+        )
+
+        rendered = template.render(
+            node=node,
+            upstream_dependencies='upstream_foo',
+            downstream_dependencies='downstream_bar'
+        )
+
+        item_name_builder_regex = re.compile(r'.*def %s_item_name_builder\(.*' % node.name)
+        batch_name_builder_regex = re.compile(r'.*def %s_batch_name_builder\(.*' % node.name)
+        filter_helper_regex = re.compile(r'.*def generator_helper_filter_with_blocklist\(.*')
+        grouped_helper_regex = re.compile(r'.*def generator_helper_grouped_list\(.*')
+        builder_invocation = r'\s+%s_builder\(\s+index = index,' % node.target
+        items_batch_name_regex = re.compile(
+            r'%s\s+items = items,\s+batch_name = batch_name,' % builder_invocation
+        )
+        item_item_name_regex = re.compile(
+            r'%s\s+item = item,\s+item_name = item_name,' % builder_invocation
+        )
+
+        return {
+            'item_name_builder': item_name_builder_regex.search(rendered),
+            'batch_name_builder': batch_name_builder_regex.search(rendered),
+            'filter_helper': filter_helper_regex.search(rendered),
+            'grouped_helper': grouped_helper_regex.search(rendered),
+            'items_batch_name': items_batch_name_regex.search(rendered),
+            'item_item_name': item_item_name_regex.search(rendered),
+        }
+
+
+HELPER = BatchingTestHelper()
+
+
+def test_batching_enabled_enabled():
+    batching_config = {'batch_size': 3}
+    node = HELPER.build_generator_node(batching_config)
+
+    assert node.batching_enabled is True
+
+
+def test_batching_enabled_disabled():
+    batching_config = {'batch_size': 3, 'disabled': True}
+    node = HELPER.build_generator_node(batching_config)
+
+    assert node.batching_enabled is False
+
+
+def test_batching_enabled_undefined():
+    node = HELPER.build_generator_node(None)
+
+    assert node.batching_enabled is False
+
+
+def test_preamble_template_batching_enabled():
+    batching_config = {'batch_size': 3}
+    matches = HELPER.run_preamble_template_test(batching_config)
+
+    assert matches['items_batch_name'] is not None
+    assert matches['item_item_name'] is None
+
+
+def test_preamble_template_batching_disabled():
+    batching_config = {'batch_size': 3, 'disabled': True}
+    matches = HELPER.run_preamble_template_test(batching_config)
+
+    assert matches['item_item_name'] is not None
+    assert matches['items_batch_name'] is None
+
+
+def test_preamble_template_batching_undefined():
+    matches = HELPER.run_preamble_template_test(None)
+
+    assert matches['item_item_name'] is not None
+    assert matches['items_batch_name'] is None
+
+
+def test_operator_template_batching_enabled():
+    batching_config = {'batch_size': 3}
+    matches = HELPER.run_operator_template_test(batching_config)
+
+    assert matches['item_name_builder'] is not None
+    assert matches['batch_name_builder'] is not None
+    assert matches['filter_helper'] is not None
+    assert matches['grouped_helper'] is not None
+    assert matches['items_batch_name'] is not None
+    assert matches['item_item_name'] is None
+
+
+def test_operator_template_batching_disabled():
+    batching_config = {'batch_size': 3, 'disabled': True}
+    matches = HELPER.run_operator_template_test(batching_config)
+
+    assert matches['item_name_builder'] is not None
+    assert matches['batch_name_builder'] is None
+    assert matches['filter_helper'] is None
+    assert matches['grouped_helper'] is None
+    assert matches['items_batch_name'] is None
+    assert matches['item_item_name'] is not None
+
+
+def test_operator_template_batching_undefined():
+    matches = HELPER.run_operator_template_test(None)
+
+    assert matches['item_name_builder'] is not None
+    assert matches['batch_name_builder'] is None
+    assert matches['filter_helper'] is None
+    assert matches['grouped_helper'] is None
+    assert matches['items_batch_name'] is None
+    assert matches['item_item_name'] is not None