Skip to content

Commit c2490a4

Browse files
authored
Fix infer list of images (#7667)
fix infer list of images
1 parent b5a2a94 commit c2490a4

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

src/datasets/arrow_writer.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from .features import Audio, Features, Image, Pdf, Value, Video
2929
from .features.features import (
3030
FeatureType,
31+
List,
3132
_ArrayXDExtensionType,
3233
_visit,
3334
cast_to_python_objects,
@@ -193,19 +194,19 @@ def _infer_custom_type_and_encode(data: Iterable) -> tuple[Iterable, Optional[Fe
193194
if isinstance(non_null_value, PIL.Image.Image):
194195
return [Image().encode_example(value) if value is not None else None for value in data], Image()
195196
if isinstance(non_null_value, list) and isinstance(non_null_value[0], PIL.Image.Image):
196-
return [[Image().encode_example(x) for x in value] if value is not None else None for value in data], [
197-
Image()
198-
]
197+
return [
198+
[Image().encode_example(x) for x in value] if value is not None else None for value in data
199+
], List(Image())
199200
if config.PDFPLUMBER_AVAILABLE and "pdfplumber" in sys.modules:
200201
import pdfplumber
201202

202203
non_null_idx, non_null_value = first_non_null_non_empty_value(data)
203204
if isinstance(non_null_value, pdfplumber.pdf.PDF):
204205
return [Pdf().encode_example(value) if value is not None else None for value in data], Pdf()
205206
if isinstance(non_null_value, list) and isinstance(non_null_value[0], pdfplumber.pdf.PDF):
206-
return [[Pdf().encode_example(x) for x in value] if value is not None else None for value in data], [
207-
Pdf()
208-
]
207+
return [
208+
[Pdf().encode_example(x) for x in value] if value is not None else None for value in data
209+
], List(Pdf())
209210
return data, None
210211

211212
def __arrow_array__(self, type: Optional[pa.DataType] = None):

0 commit comments

Comments
 (0)