From b75e1593d740c5efaecaf350405c7ffb000828a9 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Fri, 13 Dec 2024 00:43:13 -0800 Subject: [PATCH] [sonix pygrain dataset] Support FlatMap transformations. PiperOrigin-RevId: 705783333 --- grain/_src/python/dataset/transformations/flatmap.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/grain/_src/python/dataset/transformations/flatmap.py b/grain/_src/python/dataset/transformations/flatmap.py index 3941d495..885af41d 100644 --- a/grain/_src/python/dataset/transformations/flatmap.py +++ b/grain/_src/python/dataset/transformations/flatmap.py @@ -13,6 +13,7 @@ # limitations under the License. """Flatmap transformation for MapDataset.""" import functools +import sys from typing import Any, Callable, Sequence, TypeVar from grain._src.core import transforms @@ -37,6 +38,11 @@ def __init__( self._transform = transform def __len__(self) -> int: + # If the parent dataset is on infinite repeat, its length is + # sys.maxsize and would result in overflows if further increased. + # In this case, we just keep the length as sys.maxsize. + if len(self._parent) >= sys.maxsize: + return sys.maxsize return self._transform.max_fan_out * len(self._parent) def __str__(self) -> str: