fix timestamp conversion from pd to py in streaming (#4541)

Author: lhoestq

src/datasets/features/features.py

@@ -336,9 +336,23 @@ def _cast_to_python_objects(obj: Any, only_1d_for_numpy: bool, optimize_list_cas
     elif config.PIL_AVAILABLE and "PIL" in sys.modules and isinstance(obj, PIL.Image.Image):
         return encode_pil_image(obj), True
     elif isinstance(obj, pd.Series):
-        return obj.values.tolist(), True
+        return (
+            _cast_to_python_objects(
+                obj.tolist(), only_1d_for_numpy=only_1d_for_numpy, optimize_list_casting=optimize_list_casting
+            )[0],
+            True,
+        )
     elif isinstance(obj, pd.DataFrame):
-        return obj.to_dict("list"), True
+        return {
+            key: _cast_to_python_objects(
+                value, only_1d_for_numpy=only_1d_for_numpy, optimize_list_casting=optimize_list_casting
+            )[0]
+            for key, value in obj.to_dict("list").items()
+        }, True
+    elif isinstance(obj, pd.Timestamp):
+        return obj.to_pydatetime(), True
+    elif isinstance(obj, pd.Timedelta):
+        return obj.to_pytimedelta(), True
     elif isinstance(obj, Mapping):  # check for dict-like to handle nested LazyDict objects
         has_changed = not isinstance(obj, dict)
         output = {}

tests/features/test_features.py

@@ -482,6 +482,26 @@ def test_cast_to_python_objects_dataframe(self):
         casted_obj = cast_to_python_objects(obj)
         self.assertDictEqual(casted_obj, expected_obj)
 
+    def test_cast_to_python_objects_pandas_timestamp(self):
+        obj = pd.Timestamp(2020, 1, 1)
+        expected_obj = obj.to_pydatetime()
+        casted_obj = cast_to_python_objects(obj)
+        self.assertEqual(casted_obj, expected_obj)
+        casted_obj = cast_to_python_objects(pd.Series([obj]))
+        self.assertListEqual(casted_obj, [expected_obj])
+        casted_obj = cast_to_python_objects(pd.DataFrame({"a": [obj]}))
+        self.assertDictEqual(casted_obj, {"a": [expected_obj]})
+
+    def test_cast_to_python_objects_pandas_timedelta(self):
+        obj = pd.Timedelta(seconds=1)
+        expected_obj = obj.to_pytimedelta()
+        casted_obj = cast_to_python_objects(obj)
+        self.assertEqual(casted_obj, expected_obj)
+        casted_obj = cast_to_python_objects(pd.Series([obj]))
+        self.assertListEqual(casted_obj, [expected_obj])
+        casted_obj = cast_to_python_objects(pd.DataFrame({"a": [obj]}))
+        self.assertDictEqual(casted_obj, {"a": [expected_obj]})
+
     @require_torch
     def test_cast_to_python_objects_torch(self):
         import torch

tests/test_iterable_dataset.py

@@ -2,6 +2,7 @@
 from itertools import chain, islice
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from datasets import load_dataset
@@ -701,6 +702,21 @@ def test_iterable_dataset_features(features):
     assert list(dataset) == expected
 
 
+def test_iterable_dataset_features_cast_to_python():
+    ex_iterable = ExamplesIterable(
+        generate_examples_fn, {"timestamp": pd.Timestamp(2020, 1, 1), "array": np.ones(5), "n": 1}
+    )
+    features = Features(
+        {
+            "id": Value("int64"),
+            "timestamp": Value("timestamp[us]"),
+            "array": [Value("int64")],
+        }
+    )
+    dataset = IterableDataset(ex_iterable, info=DatasetInfo(features=features))
+    assert list(dataset) == [{"timestamp": pd.Timestamp(2020, 1, 1).to_pydatetime(), "array": [1] * 5, "id": 0}]
+
+
 @require_torch
 @pytest.mark.parametrize("format_type", [None, "torch", "python"])
 def test_iterable_dataset_with_format(dataset: IterableDataset, format_type):