chore: support casting floats and list-likes to timedelta series (#1362)

* chore: support casting floats and list-likes to timedelta series

* use singledispatch to handle different input types

* fix format

* stop using singledispatch due to compatibility issues

Author: sycai

bigframes/operations/timedelta_ops.py

@@ -26,6 +26,6 @@ class ToTimedeltaOp(base_ops.UnaryOp):
     unit: typing.Literal["us", "ms", "s", "m", "h", "d", "W"]
 
     def output_type(self, *input_types):
-        if input_types[0] is not dtypes.INT_DTYPE:
-            raise TypeError("expected integer input")
-        return dtypes.TIMEDELTA_DTYPE
+        if input_types[0] in (dtypes.INT_DTYPE, dtypes.FLOAT_DTYPE):
+            return dtypes.TIMEDELTA_DTYPE
+        raise TypeError("expected integer or float input")

bigframes/pandas/core/tools/timedeltas.py

@@ -18,20 +18,26 @@
     timedeltas as vendored_pandas_timedeltas,
 )
 import pandas as pd
+import pandas.api.types as pdtypes
 
 from bigframes import operations as ops
-from bigframes import series
+from bigframes import series, session
 
 
 def to_timedelta(
-    arg: typing.Union[series.Series, str, int, float],
+    arg,
     unit: typing.Optional[vendored_pandas_timedeltas.UnitChoices] = None,
-) -> typing.Union[series.Series, pd.Timedelta]:
-    if not isinstance(arg, series.Series):
-        return pd.to_timedelta(arg, unit)
+    *,
+    session: typing.Optional[session.Session] = None,
+):
+    if isinstance(arg, series.Series):
+        canonical_unit = "us" if unit is None else _canonicalize_unit(unit)
+        return arg._apply_unary_op(ops.ToTimedeltaOp(canonical_unit))
 
-    canonical_unit = "us" if unit is None else _canonicalize_unit(unit)
-    return arg._apply_unary_op(ops.ToTimedeltaOp(canonical_unit))
+    if pdtypes.is_list_like(arg):
+        return to_timedelta(series.Series(arg), unit, session=session)
+
+    return pd.to_timedelta(arg, unit)
 
 
 to_timedelta.__doc__ = vendored_pandas_timedeltas.to_timedelta.__doc__

tests/system/small/test_pandas.py

@@ -763,7 +763,7 @@ def test_to_datetime_timestamp_inputs(arg, utc, output_in_utc):
         "micros",
     ],
 )
-def test_to_timedelta_with_bf_series(session, unit):
+def test_to_timedelta_with_bf_integer_series(session, unit):
     bf_series = bpd.Series([1, 2, 3], session=session)
     pd_series = pd.Series([1, 2, 3])
 
@@ -779,6 +779,42 @@ def test_to_timedelta_with_bf_series(session, unit):
     )
 
 
+def test_to_timedelta_with_bf_float_series_value_rounded_down(session):
+    bf_series = bpd.Series([1.2, 2.9], session=session)
+
+    actual_result = (
+        typing.cast(bpd.Series, bpd.to_timedelta(bf_series, "us"))
+        .to_pandas()
+        .astype("timedelta64[ns]")
+    )
+
+    expected_result = pd.Series([pd.Timedelta(1, "us"), pd.Timedelta(2, "us")])
+    pd.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False
+    )
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        pytest.param([1, 2, 3], id="list"),
+        pytest.param((1, 2, 3), id="tuple"),
+        pytest.param(pd.Series([1, 2, 3]), id="pandas-series"),
+    ],
+)
+def test_to_timedelta_with_list_like_input(session, input):
+    actual_result = (
+        typing.cast(bpd.Series, bpd.to_timedelta(input, "s", session=session))
+        .to_pandas()
+        .astype("timedelta64[ns]")
+    )
+
+    expected_result = pd.Series(pd.to_timedelta(input, "s"))
+    pd.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False
+    )
+
+
 @pytest.mark.parametrize(
     "unit",
     ["Y", "M", "whatever"],