Handle batchable lists of length 0

Signed-off-by: eduardo apolinario <[email protected]>
flyteorg · Apr 20, 2023 · a7fc1e7 · a7fc1e7
1 parent 5e06cf1
commit a7fc1e7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
@@ -999,14 +999,17 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp
         if ListTransformer.is_batchable(python_type):
             from flytekit.types.pickle.pickle import BatchSize, FlytePickle
 
-            batchSize = len(python_val)  # default batch size
+            batch_size = len(python_val)  # default batch size
             # parse annotated to get the number of items saved in a pickle file.
             if get_origin(python_type) is Annotated:
                 for annotation in get_args(python_type)[1:]:
                     if isinstance(annotation, BatchSize):
-                        batchSize = annotation.val
+                        batch_size = annotation.val
                         break
-            lit_list = [TypeEngine.to_literal(ctx, python_val[i : i + batchSize], FlytePickle, expected.collection_type) for i in range(0, len(python_val), batchSize)]  # type: ignore
+            if batch_size > 0:
+                lit_list = [TypeEngine.to_literal(ctx, python_val[i : i + batch_size], FlytePickle, expected.collection_type) for i in range(0, len(python_val), batch_size)]  # type: ignore
+            else:
+                lit_list = []
         else:
             t = self.get_sub_type(python_type)
             lit_list = [TypeEngine.to_literal(ctx, x, t, expected.collection_type) for x in python_val]  # type: ignore

diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py
@@ -1609,6 +1609,8 @@ def test_is_batchable():
         # [[batched_FlytePickle(3 items)], [batched_FlytePickle(3 items)]]
         # Therefore, the expected list length is [2, 1] (the length of the outer list remains the same, the inner list is batched).
         ([["foo", "foo", "foo"]] * 2, typing.List[Annotated[typing.List[FlytePickle], BatchSize(3)]], [2, 1]),
+        # Case 4: Empty list
+        ([[], typing.List[FlytePickle], []]),
     ],
 )
 def test_batch_pickle_list(python_val, python_type, expected_list_length):