From d063829d6f671d5a174053566c5d0ffd41efda60 Mon Sep 17 00:00:00 2001 From: Konstantin Malanchev Date: Fri, 2 Feb 2024 16:52:05 -0500 Subject: [PATCH] Test for TsExtensionArray.__iter__ --- tests/pandas_ts/test_ts_ext_array.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/pandas_ts/test_ts_ext_array.py b/tests/pandas_ts/test_ts_ext_array.py index ec0fabb..5031791 100644 --- a/tests/pandas_ts/test_ts_ext_array.py +++ b/tests/pandas_ts/test_ts_ext_array.py @@ -138,3 +138,19 @@ def test_series_apply_udf_argument(): assert_frame_equal( series_of_dfs.iloc[0], pd.DataFrame({"a": np.array([1.0, 2.0, 3.0]), "b": -np.array([4.0, 5.0, 6.0])}) ) + + +def test___iter__(): + struct_array = pa.StructArray.from_arrays( + arrays=[ + pa.array([np.array([1.0, 2.0, 3.0]), np.array([1.0, 2.0, 1.0])]), + pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])]), + ], + names=["a", "b"], + ) + series = pd.Series(struct_array, dtype=TsDtype(struct_array.type), index=[100, 101]) + + # Check last df only + for i, df in enumerate(series): + pass + assert_frame_equal(df, pd.DataFrame({"a": np.array([1.0, 2.0, 1.0]), "b": -np.array([3.0, 4.0, 5.0])}))