chore: support timedeltas for read_pandas() (#1349)

* chore: support timedeltas for read_pandas()

* fix format

* fix mypy error

* centralize timedelta to microsecs replacement logic

* fix format

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* remove redundant imports

* polish todo comment

* update timdelta to microsecond conversion algo

* update python doc

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: sycai

bigframes/core/compile/compiler.py

@@ -24,6 +24,7 @@
 import google.cloud.bigquery
 import pandas as pd
 
+from bigframes.core import utils
 import bigframes.core.compile.compiled as compiled
 import bigframes.core.compile.concat as concat_impl
 import bigframes.core.compile.explode
@@ -173,6 +174,10 @@ def compile_readlocal(self, node: nodes.ReadLocalNode):
             io.BytesIO(node.feather_bytes),
             columns=[item.source_id for item in node.scan_list.items],
         )
+
+        # Convert timedeltas to microseconds for compatibility with BigQuery
+        _ = utils.replace_timedeltas_with_micros(array_as_pd)
+
         offsets = node.offsets_col.sql if node.offsets_col else None
         return compiled.UnorderedIR.from_pandas(
             array_as_pd, node.scan_list, offsets=offsets

bigframes/core/local_data.py

@@ -59,6 +59,9 @@ def arrow_type_replacements(type: pa.DataType) -> pa.DataType:
     if pa.types.is_time64(type):
         # This is potentially lossy, but BigFrames doesn't support ns
         return pa.time64("us")
+    if pa.types.is_duration(type):
+        # This is potentially lossy, but BigFrames doesn't support ns
+        return pa.duration("us")
     if pa.types.is_decimal128(type):
         return pa.decimal128(38, 9)
     if pa.types.is_decimal256(type):

bigframes/core/schema.py

@@ -38,9 +38,13 @@ class ArraySchema:
     items: typing.Tuple[SchemaItem, ...]
 
     @classmethod
-    def from_bq_table(cls, table: google.cloud.bigquery.Table):
+    def from_bq_table(
+        cls,
+        table: google.cloud.bigquery.Table,
+        column_type_overrides: typing.Dict[str, bigframes.dtypes.Dtype] = {},
+    ):
         items = tuple(
-            SchemaItem(name, dtype)
+            SchemaItem(name, column_type_overrides.get(name, dtype))
             for name, dtype in bigframes.dtypes.bf_type_from_type_kind(
                 table.schema
             ).items()

bigframes/core/utils.py

@@ -19,6 +19,7 @@
 
 import bigframes_vendored.pandas.io.common as vendored_pandas_io_common
 import pandas as pd
+import pandas.api.types as pdtypes
 import typing_extensions
 
 import bigframes.exceptions as bfe
@@ -184,3 +185,29 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return decorator
+
+
+def timedelta_to_micros(td: pd.Timedelta) -> int:
+    # td.value returns total nanoseconds.
+    return td.value // 1000
+
+
+def replace_timedeltas_with_micros(dataframe: pd.DataFrame) -> List[str]:
+    """
+    Replaces in-place timedeltas to integer values in microseconds. Nanosecond part is ignored.
+
+    Returns:
+        The names of updated columns
+    """
+    updated_columns = []
+
+    for col in dataframe.columns:
+        if pdtypes.is_timedelta64_dtype(dataframe[col].dtype):
+            dataframe[col] = dataframe[col].apply(timedelta_to_micros)
+            updated_columns.append(col)
+
+    if pdtypes.is_timedelta64_dtype(dataframe.index.dtype):
+        dataframe.index = dataframe.index.map(timedelta_to_micros)
+        updated_columns.append(dataframe.index.name)
+
+    return updated_columns

bigframes/dtypes.py

@@ -409,11 +409,16 @@ def dtype_for_etype(etype: ExpressionType) -> Dtype:
 def arrow_dtype_to_bigframes_dtype(arrow_dtype: pa.DataType) -> Dtype:
     if arrow_dtype in _ARROW_TO_BIGFRAMES:
         return _ARROW_TO_BIGFRAMES[arrow_dtype]
+
     if pa.types.is_list(arrow_dtype):
         return pd.ArrowDtype(arrow_dtype)
+
     if pa.types.is_struct(arrow_dtype):
         return pd.ArrowDtype(arrow_dtype)
 
+    if pa.types.is_duration(arrow_dtype):
+        return pd.ArrowDtype(arrow_dtype)
+
     # BigFrames doesn't distinguish between string and large_string because the
     # largest string (2 GB) is already larger than the largest BigQuery row.
     if pa.types.is_string(arrow_dtype) or pa.types.is_large_string(arrow_dtype):

bigframes/session/__init__.py

@@ -120,6 +120,7 @@
     pandas.ArrowDtype(pa.timestamp("us", tz="UTC")),
     pandas.ArrowDtype(pa.decimal128(38, 9)),
     pandas.ArrowDtype(pa.decimal256(76, 38)),
+    pandas.ArrowDtype(pa.duration("us")),
 )
 
 

bigframes/session/_io/pandas.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Collection, Union
+from typing import Collection, List, Union
 
 import bigframes_vendored.constants as constants
 import db_dtypes  # type: ignore
@@ -38,6 +38,7 @@ class DataFrameAndLabels:
     column_labels: Collection
     index_labels: Collection
     ordering_col: str
+    timedelta_cols: List[str]
 
 
 def _arrow_to_pandas_arrowdtype(
@@ -163,9 +164,12 @@ def pandas_to_bq_compatible(pandas_dataframe: pandas.DataFrame) -> DataFrameAndL
     pandas_dataframe_copy.columns = pandas.Index(new_col_ids)
     pandas_dataframe_copy[ordering_col] = np.arange(pandas_dataframe_copy.shape[0])
 
+    timedelta_cols = utils.replace_timedeltas_with_micros(pandas_dataframe_copy)
+
     return DataFrameAndLabels(
         df=pandas_dataframe_copy,
         column_labels=col_labels,
         index_labels=idx_labels,
         ordering_col=ordering_col,
+        timedelta_cols=timedelta_cols,
     )

bigframes/session/loader.py

@@ -176,10 +176,16 @@ def read_pandas_load_job(
         self._start_generic_job(load_job)
 
         destination_table = self._bqclient.get_table(load_table_destination)
+        col_type_overrides: typing.Dict[str, bigframes.dtypes.Dtype] = {
+            col: bigframes.dtypes.TIMEDETLA_DTYPE
+            for col in df_and_labels.timedelta_cols
+        }
         array_value = core.ArrayValue.from_table(
             table=destination_table,
-            # TODO: Generate this directly from original pandas df.
-            schema=schemata.ArraySchema.from_bq_table(destination_table),
+            # TODO (b/394156190): Generate this directly from original pandas df.
+            schema=schemata.ArraySchema.from_bq_table(
+                destination_table, col_type_overrides
+            ),
             session=self._session,
             offsets_col=ordering_col,
         ).drop_columns([ordering_col])
@@ -229,10 +235,16 @@ def read_pandas_streaming(
                     f"Problem loading at least one row from DataFrame: {errors}. {constants.FEEDBACK_LINK}"
                 )
 
+        col_type_overrides: typing.Dict[str, bigframes.dtypes.Dtype] = {
+            col: bigframes.dtypes.TIMEDETLA_DTYPE
+            for col in df_and_labels.timedelta_cols
+        }
         array_value = (
             core.ArrayValue.from_table(
                 table=destination_table,
-                schema=schemata.ArraySchema.from_bq_table(destination_table),
+                schema=schemata.ArraySchema.from_bq_table(
+                    destination_table, col_type_overrides
+                ),
                 session=self._session,
                 # Don't set the offsets column because we want to group by it.
             )

tests/system/small/test_session.py

@@ -691,6 +691,62 @@ def test_read_pandas_tokyo(
     assert len(expected) == result.total_rows
 
 
+@pytest.mark.parametrize(
+    "write_engine",
+    ["default", "bigquery_inline", "bigquery_load", "bigquery_streaming"],
+)
+def test_read_pandas_timedelta_dataframes(session, write_engine):
+    expected_df = pd.DataFrame({"my_col": pd.to_timedelta([1, 2, 3], unit="d")})
+
+    actual_result = (
+        session.read_pandas(expected_df, write_engine=write_engine)
+        .to_pandas()
+        .astype("timedelta64[ns]")
+    )
+
+    if write_engine == "bigquery_streaming":
+        expected_df.index = pd.Index([pd.NA] * 3, dtype="Int64")
+    pd.testing.assert_frame_equal(actual_result, expected_df, check_index_type=False)
+
+
+@pytest.mark.parametrize(
+    "write_engine",
+    ["default", "bigquery_inline", "bigquery_load", "bigquery_streaming"],
+)
+def test_read_pandas_timedelta_series(session, write_engine):
+    expected_series = pd.Series(pd.to_timedelta([1, 2, 3], unit="d"))
+
+    actual_result = (
+        session.read_pandas(expected_series, write_engine=write_engine)
+        .to_pandas()
+        .astype("timedelta64[ns]")
+    )
+
+    if write_engine == "bigquery_streaming":
+        expected_series.index = pd.Index([pd.NA] * 3, dtype="Int64")
+    pd.testing.assert_series_equal(
+        actual_result, expected_series, check_index_type=False
+    )
+
+
+@pytest.mark.parametrize(
+    "write_engine",
+    ["default", "bigquery_inline", "bigquery_load"],
+)
+def test_read_pandas_timedelta_index(session, write_engine):
+    expected_index = pd.to_timedelta(
+        [1, 2, 3], unit="d"
+    )  # to_timedelta returns an index
+
+    actual_result = (
+        session.read_pandas(expected_index, write_engine=write_engine)
+        .to_pandas()
+        .astype("timedelta64[ns]")
+    )
+
+    pd.testing.assert_index_equal(actual_result, expected_index)
+
+
 @utils.skip_legacy_pandas
 @pytest.mark.parametrize(
     ("write_engine",),