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,10 @@ class GeneratorNode(SubdagNode):
     def regex_blocklist(self):
         return self.item.get('regex_blocklist', ())
 
+    @property
+    def batching(self):
+        return self.item.get('batching', {'enabled': False, 'batch_size': 1})
+
 
 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_config = execution_context.referrer.item.get('batching', {'enabled': False})
+        # Validate suffix_mode based on batching config
+        if batching_config['enabled'] and suffix_mode == 'item_name':
+            raise Exception(
+                'Cannot use `item_name` for auto_task_id_mode when batching is enabled')
+        elif not batching_config['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_config['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

@@ -14,7 +14,7 @@
 #     limitations under the License.
 
 import semver
-from marshmallow import fields, validates_schema, ValidationError
+from marshmallow import fields, post_load, pre_dump, validates_schema, ValidationError
 from boundary_layer import VERSION, MIN_SUPPORTED_VERSION
 from boundary_layer.schemas.base import StrictSchema
 
@@ -35,9 +35,35 @@ class ReferenceSchema(OperatorSchema):
     target = fields.String(required=True)
 
 
+class BatchingSchema(StrictSchema):
+    enabled = fields.Boolean()
+    batch_size = fields.Integer(required=True)
+    # This is a "transient" field to help with implicit enablement behavior
+    original_enabled = fields.Boolean(load_only=True)
+
+    @post_load
+    def fix_enabled_pre_load(self, data):
+        """
+        If batching config is set at all, it's assumed to be enabled.
+        """
+        enabled = data.get('enabled', None)
+        data['original_enabled'] = enabled
+        if enabled is None:
+            data['enabled'] = True
+
+    @pre_dump
+    def fix_enabled_pre_dump(self, data):
+        """
+        Don't persist the enabled field if it wasn't explicitly configured.
+        """
+        if data['original_enabled'] is None:
+            del data['enabled']
+
+
 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_schemas.py

@@ -0,0 +1,57 @@
+from boundary_layer.schemas.dag import BatchingSchema
+
+
+def test_batching_schema_implicit_enabled():
+    schema = BatchingSchema()
+    data = {
+        'batch_size': 10
+    }
+    batching = schema.load(data)[0]
+
+    assert batching['enabled'] is True
+    assert batching['batch_size'] == 10
+    assert batching['original_enabled'] is None
+
+    dumped = schema.dump(batching)[0]
+
+    assert 'enabled' not in dumped
+    assert dumped['batch_size'] == 10
+    assert 'original_enabled' not in dumped
+
+
+def test_batching_schema_explicit_enabled():
+    schema = BatchingSchema()
+    data = {
+        'enabled': True,
+        'batch_size': 10
+    }
+    batching = schema.load(data)[0]
+
+    assert batching['enabled'] is True
+    assert batching['batch_size'] == 10
+    assert batching['original_enabled'] is True
+
+    dumped = schema.dump(batching)[0]
+
+    assert dumped['enabled'] is True
+    assert dumped['batch_size'] == 10
+    assert 'original_enabled' not in dumped
+
+
+def test_batching_schema_disabled():
+    schema = BatchingSchema()
+    data = {
+        'enabled': False,
+        'batch_size': 10
+    }
+    batching = schema.load(data)[0]
+
+    assert batching['enabled'] is False
+    assert batching['batch_size'] == 10
+    assert batching['original_enabled'] is False
+
+    dumped = schema.dump(batching)[0]
+
+    assert dumped['enabled'] is False
+    assert dumped['batch_size'] == 10
+    assert 'original_enabled' not in dumped