Test for TsExtensionArray.__iter__

lincc-frameworks · Feb 2, 2024 · d063829 · d063829
1 parent 7c612c4
commit d063829
Showing 1 changed file with 16 additions and 0 deletions.
diff --git 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])}))