From c35e82962f3fa3caf4f2337e4df009a6d76bf0d8 Mon Sep 17 00:00:00 2001 From: Mike Boss Date: Fri, 5 Apr 2024 15:19:24 +0200 Subject: [PATCH] get_inferred_type calls pa.array(self) which is the same as pa.array(types_sequence) --- src/datasets/arrow_writer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/datasets/arrow_writer.py b/src/datasets/arrow_writer.py index 82e72a91ecc..08fb680ce65 100644 --- a/src/datasets/arrow_writer.py +++ b/src/datasets/arrow_writer.py @@ -561,8 +561,9 @@ def write_batch( else: col_try_type = try_features[col] if try_features is not None and col in try_features else None typed_sequence = OptimizedTypedSequence(col_values, type=col_type, try_type=col_try_type, col=col) - arrays.append(pa.array(typed_sequence)) - inferred_features[col] = typed_sequence.get_inferred_type() + array = pa.array(typed_sequence) + arrays.append(array) + inferred_features[col] = generate_from_arrow_type(array.type) schema = inferred_features.arrow_schema if self.pa_writer is None else self.schema pa_table = pa.Table.from_arrays(arrays, schema=schema) self.write_table(pa_table, writer_batch_size)