Partition limit (#2301)

Add a partition limit. Signed-off-by: Yee Hing Tong <[email protected]> Signed-off-by: Jan Fiedler <[email protected]>
flyteorg · Jul 25, 2024 · 59ac755 · 59ac755
1 parent eb294c9
commit 59ac755
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/flytekit/core/artifact.py b/flytekit/core/artifact.py
@@ -15,6 +15,7 @@
 from flytekit.loggers import logger
 
 TIME_PARTITION_KWARG = "time_partition"
+MAX_PARTITIONS = 10
 
 
 class InputsBase(object):
@@ -337,6 +338,9 @@ def __init__(
             self._partitions = Partitions(p)
             self._partitions.set_reference_artifact(self)
 
+        if self.partition_keys and len(self.partition_keys) > MAX_PARTITIONS:
+            raise ValueError("There is a hard limit of 10 partition keys per artifact currently.")
+
     def __call__(self, *args, **kwargs) -> ArtifactIDSpecification:
         """
         This __call__ should only ever happen in the context of a task or workflow's output, to be

diff --git a/tests/flytekit/unit/core/test_artifacts.py b/tests/flytekit/unit/core/test_artifacts.py
@@ -583,3 +583,9 @@ def test_tp_math():
     assert tp2.other == datetime.timedelta(days=1)
     assert tp2.granularity == Granularity.HOUR
     assert tp2 is not tp
+
+
+def test_lims():
+    # test an artifact with 11 partition keys
+    with pytest.raises(ValueError):
+        Artifact(name="test artifact", time_partitioned=True, partition_keys=[f"key_{i}" for i in range(11)])