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 #}
+def 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 #}
+{# https://stackoverflow.com/a/312464 #}
+def grouped(l, n):
+    for i in range(0, len(l), n):
+        yield l[i:i + n]
+
+{% set filtered = (node.name + '_filtered') | sanitize_operator_name %}
+{{ filtered }} = filter_with_blocklist({{ all_items }}, {{ item_name_builder }}, {{ blocklist }})
+
+for (index, items) in enumerate(grouped({{ 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,9 @@ 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>>'
+        name_var = 'batch_name' if execution_context.referrer.item.get('batching', {'enabled': False})['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):
+    enabled = fields.Boolean(required=True)
+    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):