Merge pull request #291 from mlcommons/feature/refacto-group-record-set

Only create a `Join` if it is needed.
mlcommons · Nov 3, 2023 · cfa6ba3 · cfa6ba3
2 parents d5dfda7 + 5fae274
commit cfa6ba3
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/mlcroissant/mlcroissant/_src/operation_graph/graph.py b/python/mlcroissant/mlcroissant/_src/operation_graph/graph.py
@@ -52,9 +52,12 @@ def _add_operations_for_field_with_source(
     - `GroupRecordSetStart` to structure the final dict that is sent back to the user.
     """
     record_set = _find_record_set(node)
+    operation = operations.last_operations(node, only_leaf=True)
+    has_join = any(field for field in record_set.fields if field.references.uid)
+    if has_join:
+        operation = [operation >> Join(operations=operations, node=record_set)]
     (
-        operations.last_operations(node, only_leaf=True)
-        >> Join(operations=operations, node=record_set)
+        operation
         >> GroupRecordSetStart(operations=operations, node=record_set)
         >> ReadField(operations=operations, node=node)
         >> GroupRecordSetEnd(operations=operations, node=record_set)