chore: add type checks for unix epoch conversions (#1343)

Author: sycai

bigframes/operations/datetime_ops.py

@@ -43,7 +43,7 @@ class ToDatetimeOp(base_ops.UnaryOp):
     format: typing.Optional[str] = None
     unit: typing.Optional[str] = None
 
-    def output_type(self, *input_types):
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         if input_types[0] not in (
             dtypes.FLOAT_DTYPE,
             dtypes.INT_DTYPE,
@@ -59,7 +59,7 @@ class ToTimestampOp(base_ops.UnaryOp):
     format: typing.Optional[str] = None
     unit: typing.Optional[str] = None
 
-    def output_type(self, *input_types):
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         # Must be numeric or string
         if input_types[0] not in (
             dtypes.FLOAT_DTYPE,
@@ -75,29 +75,35 @@ class StrftimeOp(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "strftime"
     date_format: str
 
-    def output_type(self, *input_types):
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.STRING_DTYPE
 
 
 @dataclasses.dataclass(frozen=True)
 class UnixSeconds(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_seconds"
 
-    def output_type(self, *input_types):
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+            raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
 
 @dataclasses.dataclass(frozen=True)
 class UnixMillis(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_millis"
 
-    def output_type(self, *input_types):
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+            raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
 
 @dataclasses.dataclass(frozen=True)
 class UnixMicros(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_micros"
 
-    def output_type(self, *input_types):
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+            raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE

tests/system/small/bigquery/test_datetime.py

@@ -15,6 +15,7 @@
 import typing
 
 import pandas as pd
+import pytest
 
 from bigframes import bigquery
 
@@ -32,6 +33,13 @@ def test_unix_seconds(scalars_dfs):
     pd.testing.assert_series_equal(actual_res, expected_res)
 
 
+def test_unix_seconds_incorrect_input_type_raise_error(scalars_dfs):
+    df, _ = scalars_dfs
+
+    with pytest.raises(TypeError):
+        bigquery.unix_seconds(df["string_col"])
+
+
 def test_unix_millis(scalars_dfs):
     bigframes_df, pandas_df = scalars_dfs
 
@@ -45,6 +53,13 @@ def test_unix_millis(scalars_dfs):
     pd.testing.assert_series_equal(actual_res, expected_res)
 
 
+def test_unix_millis_incorrect_input_type_raise_error(scalars_dfs):
+    df, _ = scalars_dfs
+
+    with pytest.raises(TypeError):
+        bigquery.unix_millis(df["string_col"])
+
+
 def test_unix_micros(scalars_dfs):
     bigframes_df, pandas_df = scalars_dfs
 
@@ -58,6 +73,13 @@ def test_unix_micros(scalars_dfs):
     pd.testing.assert_series_equal(actual_res, expected_res)
 
 
+def test_unix_micros_incorrect_input_type_raise_error(scalars_dfs):
+    df, _ = scalars_dfs
+
+    with pytest.raises(TypeError):
+        bigquery.unix_micros(df["string_col"])
+
+
 def _to_unix_epoch(
     ts: pd.Timestamp, unit: typing.Literal["s", "ms", "us"]
 ) -> typing.Optional[int]: