[AL-4870] Split model run data rows using global keys

Author: whistler

CHANGELOG.md

@@ -2,8 +2,9 @@
 
 # Version 3.40.0 (YYYY-MM-DD)
 
-## Added 
-* Insert newest changelogs here
+## Added
+* Support Global keys to reference data rows in `Project.create_batch()`, `ModelRun.assign_data_rows_to_split()`.
+
 
 # Version 3.39.0 (2023-02-28)
 ## Added

labelbox/schema/model_run.py

@@ -335,14 +335,15 @@ def delete_model_run_data_rows(self, data_row_ids: List[str]):
 
     @experimental
     def assign_data_rows_to_split(self,
-                                  data_row_ids: List[str],
-                                  split: Union[DataSplit, str],
+                                  data_row_ids: List[str] = None,
+                                  split: Union[DataSplit, str] = None,
+                                  global_keys: List[str] = None,
                                   timeout_seconds=120):
 
         split_value = split.value if isinstance(split, DataSplit) else split
         valid_splits = DataSplit._member_names_
 
-        if split_value not in valid_splits:
+        if split_value is None or split_value not in valid_splits:
             raise ValueError(
                 f"`split` must be one of : `{valid_splits}`. Found : `{split}`")
 
@@ -354,7 +355,8 @@ def assign_data_rows_to_split(self,
                 'data': {
                     'assignments': [{
                         'split': split_value,
-                        'dataRowIds': data_row_ids
+                        'dataRowIds': data_row_ids,
+                        'globalKeys': global_keys,
                     }]
                 }
             },

tests/integration/annotation_import/test_model_run.py

@@ -1,5 +1,6 @@
 import time
 import os
+import uuid
 import pytest
 
 from collections import Counter
@@ -208,7 +209,8 @@ def test_model_run_export_v2(model_run_with_model_run_data_rows,
             assert prediction_id in label_ids_set
 
 
-def test_model_run_split_assignment(model_run, dataset, image_url):
+def test_model_run_split_assignment_by_data_row_ids(model_run, dataset,
+                                                    image_url):
     n_data_rows = 10
     data_rows = dataset.create_data_rows([{
         "row_data": image_url
@@ -227,3 +229,18 @@ def test_model_run_split_assignment(model_run, dataset, image_url):
             counts[data_row.data_split.value] += 1
         split = split.value if isinstance(split, DataSplit) else split
         assert counts[split] == n_data_rows
+
+
+def test_model_run_split_assignment_by_global_keys(model_run, data_rows):
+    global_keys = [data_row.global_key for data_row in data_rows]
+
+    model_run.upsert_data_rows(global_keys=global_keys)
+
+    for split in ["TRAINING", "TEST", "VALIDATION", "UNASSIGNED", *DataSplit]:
+        model_run.assign_data_rows_to_split(split=split,
+                                            global_keys=global_keys)
+        splits = [
+            data_row.data_split.value
+            for data_row in model_run.model_run_data_rows()
+        ]
+        assert len(set(splits)) == 1