wip

lib/repack/owid/repack/__init__.py

@@ -77,7 +77,8 @@ def to_int(s: pd.Series) -> pd.Series:
     # values could be integers or strings
     v = s.astype("float64").astype("int64[pyarrow]")
 
-    if not series_eq(v, s, cast="float64[pyarrow]"):
+    # casting to float converts strings to floats, that doesn't work with float64[pyarrow]
+    if not series_eq(v, s, cast=float):
         raise ValueError()
 
     # it's an integer, now pack it smaller
@@ -145,7 +146,7 @@ def series_eq(lhs: pd.Series, rhs: pd.Series, cast: Any, rtol: float = 1e-5, ato
         return False
 
     # improve performance by calling native astype method
-    if cast == "float64[pyarrow]":
+    if cast in {float, "float", "float64[pyarrow]"}:
         func = lambda s: s.astype(cast)  # noqa: E731
     else:
         raise NotImplementedError()

lib/repack/tests/test_repack.py

@@ -27,6 +27,7 @@ def test_repack_object_columns():
             "myint": [1, 2, None, 3],
             "myfloat": [1.2, 2.0, 3.0, None],
             "mycat": ["a", None, "b", "c"],
+            "myintstr": [1, 2, 3, "4"],
         },
         dtype="object",
     )
@@ -37,6 +38,7 @@ def test_repack_object_columns():
     assert df_repack.myint.dtype.name == "UInt8"
     assert df_repack.myfloat.dtype.name == "float32"
     assert df_repack.mycat.dtype.name == "category"
+    assert df_repack.myintstr.dtype.name == "UInt8"
 
 
 def test_repack_frame_with_index():