feat(snowflake): implement extract millisecond

ibis-project · Jan 19, 2023 · 3292e91 · 3292e91
1 parent 271554c
commit 3292e91
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 10 deletions.
diff --git a/ibis/backends/snowflake/registry.py b/ibis/backends/snowflake/registry.py
@@ -227,6 +227,12 @@ def _map(_, op):
         ops.RegexExtract: fixed_arity(sa.func.regexp_substr, 3),
         ops.RegexSearch: fixed_arity(lambda left, right: left.op('REGEXP')(right), 2),
         ops.RegexReplace: fixed_arity(sa.func.regexp_replace, 3),
+        ops.ExtractMillisecond: fixed_arity(
+            lambda arg: sa.cast(
+                sa.extract("epoch_millisecond", arg) % 1000, sa.SMALLINT
+            ),
+            1,
+        ),
     }
 )
 
@@ -247,7 +253,6 @@ def _map(_, op):
     # ibis.expr.operations.structs
     ops.StructField,
     # ibis.expr.operations.temporal
-    ops.ExtractMillisecond,
     ops.IntervalFromInteger,
     ops.TimestampDiff,
     ops.TimestampFromUNIX,

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -77,13 +77,7 @@ def test_timestamp_extract(backend, alltypes, df, attr):
             methodcaller('millisecond'),
             359,
             id='millisecond',
-            marks=[
-                pytest.mark.notimpl(["clickhouse", "pyspark"]),
-                pytest.mark.notyet(
-                    ["snowflake"],
-                    reason="milliseconds doesn't seem to work at all upstream",
-                ),
-            ],
+            marks=[pytest.mark.notimpl(["clickhouse", "pyspark"])],
         ),
         param(
             lambda x: x.day_of_week.index(),
@@ -104,7 +98,7 @@ def test_timestamp_extract_literal(con, func, expected):
     assert con.execute(func(value).name("tmp")) == expected
 
 
-@pytest.mark.notimpl(["datafusion", "clickhouse", "snowflake"])
+@pytest.mark.notimpl(["datafusion", "clickhouse"])
 @pytest.mark.notyet(["sqlite", "pyspark"])
 def test_timestamp_extract_milliseconds(backend, alltypes, df):
     expr = alltypes.timestamp_col.millisecond().name("millisecond")
@@ -855,7 +849,7 @@ def test_date_column_from_iso(con, alltypes, df):
     tm.assert_series_equal(golden.rename('tmp'), actual.rename('tmp'))
 
 
-@pytest.mark.notimpl(["datafusion", "snowflake"])
+@pytest.mark.notimpl(["datafusion"])
 @pytest.mark.notyet(["clickhouse", "pyspark"])
 def test_timestamp_extract_milliseconds_with_big_value(con):
     timestamp = ibis.timestamp("2021-01-01 01:30:59.333456")