diff --git a/ibis/backends/snowflake/registry.py b/ibis/backends/snowflake/registry.py index fd37e2d5c667..31e7e71a0d79 100644 --- a/ibis/backends/snowflake/registry.py +++ b/ibis/backends/snowflake/registry.py @@ -257,6 +257,24 @@ def _map_get(t, op): return sa.cast(expr, sqla_type) +def _timestamp_bucket(t, op): + if op.offset is not None: + raise com.UnsupportedOperationError( + "`offset` is not supported in the Snowflake backend for timestamp bucketing" + ) + + interval = op.interval + + if not isinstance(interval, ops.Literal): + raise com.UnsupportedOperationError( + f"Interval must be a literal for the Snowflake backend, got {type(interval)}" + ) + + return sa.func.time_slice( + t.translate(op.arg), interval.value, interval.dtype.unit.name + ) + + _TIMESTAMP_UNITS_TO_SCALE = {"s": 0, "ms": 3, "us": 6, "ns": 9} _SF_POS_INF = sa.func.to_double("Inf") @@ -471,6 +489,7 @@ def _map_get(t, op): ops.TimestampDelta: fixed_arity( lambda part, left, right: sa.func.timestampdiff(part, right, left), 3 ), + ops.TimestampBucket: _timestamp_bucket, } ) diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py index 99ab454529dc..e1bee212d853 100644 --- a/ibis/backends/tests/test_temporal.py +++ b/ibis/backends/tests/test_temporal.py @@ -2526,7 +2526,6 @@ def test_delta(con, start, end, unit, expected): "oracle", "pandas", "pyspark", - "snowflake", "sqlite", "trino", ], @@ -2543,16 +2542,24 @@ def test_delta(con, start, end, unit, expected): param( {"milliseconds": 50}, "50ms", - marks=pytest.mark.notimpl( - ["clickhouse"], - raises=com.UnsupportedOperationError, - reason="clickhouse doesn't support sub-second interval precision", - ), + marks=[ + pytest.mark.notimpl( + ["clickhouse"], + raises=com.UnsupportedOperationError, + reason="backend doesn't support sub-second interval precision", + ), + pytest.mark.notimpl( + ["snowflake"], + raises=sa.exc.ProgrammingError, + reason="snowflake doesn't support sub-second interval precision", + ), + ], + id="milliseconds", ), - ({"seconds": 2}, "2s"), - ({"minutes": 5}, "300s"), - ({"hours": 2}, "2h"), - ({"days": 2}, "2D"), + param({"seconds": 2}, "2s", id="seconds"), + param({"minutes": 5}, "300s", id="minutes"), + param({"hours": 2}, "2h", id="hours"), + param({"days": 2}, "2D", id="days"), ], ) def test_timestamp_bucket(backend, kws, pd_freq): @@ -2573,7 +2580,6 @@ def test_timestamp_bucket(backend, kws, pd_freq): "oracle", "pandas", "pyspark", - "snowflake", "sqlite", "trino", ], @@ -2585,11 +2591,11 @@ def test_timestamp_bucket(backend, kws, pd_freq): reason="Druid tests load timestamp_col as a string currently", ) @pytest.mark.notimpl( - ["clickhouse", "mssql"], + ["clickhouse", "mssql", "snowflake"], reason="offset arg not supported", raises=com.UnsupportedOperationError, ) -@pytest.mark.parametrize("offset_mins", [2, -2]) +@pytest.mark.parametrize("offset_mins", [2, -2], ids=["pos", "neg"]) def test_timestamp_bucket_offset(backend, offset_mins): ts = backend.functional_alltypes.timestamp_col.name("ts") expr = ts.bucket(minutes=5, offset=ibis.interval(minutes=offset_mins)).name("ts")