fix(pandas): fix integer wraparound when extracting epoch seconds

ibis-project · Dec 18, 2023 · e98fa3c · e98fa3c
1 parent 7bb0470
commit e98fa3c
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 6 deletions.
diff --git a/ibis/backends/dask/execution/temporal.py b/ibis/backends/dask/execution/temporal.py
@@ -32,7 +32,7 @@
     execute_date_sub_diff_series_date,
     execute_day_of_week_index_series,
     execute_day_of_week_name_series,
-    execute_epoch_seconds,
+    execute_epoch_seconds_series,
     execute_extract_microsecond_series,
     execute_extract_millisecond_series,
     execute_extract_timestamp_field_series,
@@ -61,7 +61,7 @@
     ops.ExtractTemporalField: [((dd.Series,), execute_extract_timestamp_field_series)],
     ops.ExtractMicrosecond: [((dd.Series,), execute_extract_microsecond_series)],
     ops.ExtractMillisecond: [((dd.Series,), execute_extract_millisecond_series)],
-    ops.ExtractEpochSeconds: [((dd.Series,), execute_epoch_seconds)],
+    ops.ExtractEpochSeconds: [((dd.Series,), execute_epoch_seconds_series)],
     ops.IntervalFromInteger: [((dd.Series,), execute_interval_from_integer_series)],
     ops.IntervalAdd: [
         (

diff --git a/ibis/backends/pandas/execution/temporal.py b/ibis/backends/pandas/execution/temporal.py
@@ -65,9 +65,19 @@ def execute_extract_microsecond_series(op, data, **kwargs):
     return data.dt.microsecond.astype(np.int32)
 
 
-@execute_node.register(ops.ExtractEpochSeconds, (datetime.datetime, pd.Series))
-def execute_epoch_seconds(op, data, **kwargs):
-    return data.astype("datetime64[s]").astype("int64").astype("int32")
+@execute_node.register(ops.ExtractEpochSeconds, pd.Series)
+def execute_epoch_seconds_series(op, data, **kwargs):
+    return (
+        data.astype("datetime64[ns]")
+        .astype("int64")
+        .floordiv(1_000_000_000)
+        .astype("int32")
+    )
+
+
+@execute_node.register(ops.ExtractEpochSeconds, (pd.Timestamp, datetime.datetime))
+def execute_epoch_seconds_literal(op, data, **kwargs):
+    return pd.Timestamp(data).floor("s").value // 1_000_000_000
 
 
 @execute_node.register(

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -354,7 +354,11 @@ def test_timestamp_extract_epoch_seconds(backend, alltypes, df):
     result = expr.execute()
 
     expected = backend.default_series_rename(
-        df.timestamp_col.astype("datetime64[s]").astype("int64").astype("int32")
+        df.timestamp_col.astype("datetime64[ns]")
+        .dt.floor("s")
+        .astype("int64")
+        .floordiv(1_000_000_000)
+        .astype("int32")
     )
     backend.assert_series_equal(result, expected)