From 18b212be913cb768a6fcd3a6de8754ce727650ac Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Thu, 20 Apr 2023 10:11:54 -0700 Subject: [PATCH] Handle batchable lists of length 0 (#1600) Signed-off-by: eduardo apolinario Co-authored-by: eduardo apolinario --- flytekit/core/type_engine.py | 9 ++++++--- tests/flytekit/unit/core/test_type_engine.py | 2 ++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 009d2b1e7b..5672d9ed06 100644 --- 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 index d70174c77a..ada30683d8 100644 --- 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):