fix: use value equality to check types for unix epoch functions and timestamp diff

bigframes/operations/datetime_ops.py

@@ -84,7 +84,7 @@ class UnixSeconds(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_seconds"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+        if input_types[0] != dtypes.TIMESTAMP_DTYPE:
             raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
@@ -94,7 +94,7 @@ class UnixMillis(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_millis"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+        if input_types[0] != dtypes.TIMESTAMP_DTYPE:
             raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
@@ -104,7 +104,7 @@ class UnixMicros(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_micros"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+        if input_types[0] != dtypes.TIMESTAMP_DTYPE:
             raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
@@ -114,7 +114,7 @@ class TimestampDiff(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "timestamp_diff"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not input_types[1]:
+        if input_types[0] != input_types[1]:
             raise TypeError(
                 f"two inputs have different types. left: {input_types[0]}, right: {input_types[1]}"
             )

tests/system/small/bigquery/test_datetime.py

@@ -15,10 +15,20 @@
 import typing
 
 import pandas as pd
+import pyarrow as pa
 import pytest
 
 from bigframes import bigquery
 
+_TIMESTAMP_DTYPE = pd.ArrowDtype(pa.timestamp("us", tz="UTC"))
+
+
+@pytest.fixture
+def int_series(session):
+    pd_series = pd.Series([1, 2, 3, 4, 5])
+
+    return session.read_pandas(pd_series), pd_series
+
 
 def test_unix_seconds(scalars_dfs):
     bigframes_df, pandas_df = scalars_dfs
@@ -33,6 +43,19 @@ def test_unix_seconds(scalars_dfs):
     pd.testing.assert_series_equal(actual_res, expected_res)
 
 
+def test_unix_seconds_after_type_casting(int_series):
+    bf_series, pd_series = int_series
+
+    actual_res = bigquery.unix_seconds(bf_series.astype(_TIMESTAMP_DTYPE)).to_pandas()
+
+    expected_res = (
+        pd_series.astype(_TIMESTAMP_DTYPE)
+        .apply(lambda ts: _to_unix_epoch(ts, "s"))
+        .astype("Int64")
+    )
+    pd.testing.assert_series_equal(actual_res, expected_res, check_index_type=False)
+
+
 def test_unix_seconds_incorrect_input_type_raise_error(scalars_dfs):
     df, _ = scalars_dfs
 
@@ -53,6 +76,19 @@ def test_unix_millis(scalars_dfs):
     pd.testing.assert_series_equal(actual_res, expected_res)
 
 
+def test_unix_millis_after_type_casting(int_series):
+    bf_series, pd_series = int_series
+
+    actual_res = bigquery.unix_millis(bf_series.astype(_TIMESTAMP_DTYPE)).to_pandas()
+
+    expected_res = (
+        pd_series.astype(_TIMESTAMP_DTYPE)
+        .apply(lambda ts: _to_unix_epoch(ts, "ms"))
+        .astype("Int64")
+    )
+    pd.testing.assert_series_equal(actual_res, expected_res, check_index_type=False)
+
+
 def test_unix_millis_incorrect_input_type_raise_error(scalars_dfs):
     df, _ = scalars_dfs
 
@@ -73,6 +109,19 @@ def test_unix_micros(scalars_dfs):
     pd.testing.assert_series_equal(actual_res, expected_res)
 
 
+def test_unix_micros_after_type_casting(int_series):
+    bf_series, pd_series = int_series
+
+    actual_res = bigquery.unix_micros(bf_series.astype(_TIMESTAMP_DTYPE)).to_pandas()
+
+    expected_res = (
+        pd_series.astype(_TIMESTAMP_DTYPE)
+        .apply(lambda ts: _to_unix_epoch(ts, "us"))
+        .astype("Int64")
+    )
+    pd.testing.assert_series_equal(actual_res, expected_res, check_index_type=False)
+
+
 def test_unix_micros_incorrect_input_type_raise_error(scalars_dfs):
     df, _ = scalars_dfs
 

tests/system/small/operations/test_timedeltas.py

@@ -60,6 +60,7 @@ def temporal_dfs(session):
             ],
             "float_col": [1.5, 2, -3],
             "int_col": [1, 2, -3],
+            "positive_int_col": [1, 2, 3],
         }
     )
 
@@ -607,3 +608,24 @@ def test_timedelta_agg__int_result(temporal_dfs, agg_func):
 
     expected_result = agg_func(pd_df["timedelta_col_1"])
     assert actual_result == expected_result
+
+
+def test_timestamp_diff_after_type_casting(temporal_dfs):
+    if version.Version(pd.__version__) <= version.Version("2.1.0"):
+        pytest.skip(
+            "Temporal type casting is not well-supported in older verions of Pandas."
+        )
+
+    bf_df, pd_df = temporal_dfs
+    dtype = pd.ArrowDtype(pa.timestamp("us", tz="UTC"))
+
+    actual_result = (
+        bf_df["timestamp_col"] - bf_df["positive_int_col"].astype(dtype)
+    ).to_pandas()
+
+    expected_result = pd_df["timestamp_col"] - pd_df["positive_int_col"].astype(
+        "datetime64[us, UTC]"
+    )
+    pandas.testing.assert_series_equal(
+        actual_result, expected_result, check_index_type=False, check_dtype=False
+    )