From 1ac71f788d85153128451d706877c765390c47e7 Mon Sep 17 00:00:00 2001 From: mariosasko Date: Mon, 4 Jul 2022 18:07:50 +0200 Subject: [PATCH] Fix time type _arrow_to_datasets_dtype conversion --- src/datasets/features/features.py | 4 ++-- tests/features/test_features.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/datasets/features/features.py b/src/datasets/features/features.py index 71a63a0f9d2..4109e639c67 100644 --- a/src/datasets/features/features.py +++ b/src/datasets/features/features.py @@ -78,9 +78,9 @@ def _arrow_to_datasets_dtype(arrow_type: pa.DataType) -> str: elif pyarrow.types.is_float64(arrow_type): return "float64" # pyarrow dtype is "double" elif pyarrow.types.is_time32(arrow_type): - return f"time32[{arrow_type.unit}]" + return f"time32[{pa.type_for_alias(str(arrow_type)).unit}]" elif pyarrow.types.is_time64(arrow_type): - return f"time64[{arrow_type.unit}]" + return f"time64[{pa.type_for_alias(str(arrow_type)).unit}]" elif pyarrow.types.is_timestamp(arrow_type): if arrow_type.tz is None: return f"timestamp[{arrow_type.unit}]" diff --git a/tests/features/test_features.py b/tests/features/test_features.py index b22807b56c2..cc79442b231 100644 --- a/tests/features/test_features.py +++ b/tests/features/test_features.py @@ -1,3 +1,4 @@ +import datetime from dataclasses import asdict from unittest import TestCase from unittest.mock import patch @@ -58,6 +59,7 @@ def test_string_to_arrow_bijection_for_primitive_types(self): pa.string(), pa.int32(), pa.float64(), + pa.array([datetime.time(1, 1, 1)]).type, # arrow type: DataType(time64[us]) ] for dt in supported_pyarrow_datatypes: self.assertEqual(dt, string_to_arrow(_arrow_to_datasets_dtype(dt)))