feat(flink): implement the ops.TryCast operation

ibis-project · Oct 3, 2023 · 752e587 · 752e587
1 parent 5cd9082
commit 752e587
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/ibis/backends/flink/registry.py b/ibis/backends/flink/registry.py
@@ -40,6 +40,18 @@ def _literal(translator: ExprTranslator, op: ops.Literal) -> str:
     return translate_literal(op)
 
 
+def _try_cast(translator: ExprTranslator, op: ops.Node) -> str:
+    arg_formatted = translator.translate(op.arg)
+
+    if op.arg.dtype.is_temporal() and op.to.is_numeric():
+        # The cast from TIMESTAMP type to NUMERIC type is not allowed.
+        # It's recommended to use UNIX_TIMESTAMP(CAST(timestamp_col AS STRING)) instead.
+        return f"UNIX_TIMESTAMP(TRY_CAST({arg_formatted} AS STRING))"
+    else:
+        sql_type = helpers.type_to_sql_string(op.to)
+        return f"TRY_CAST({arg_formatted} AS {sql_type})"
+
+
 def _filter(translator: ExprTranslator, op: ops.Node) -> str:
     bool_expr = translator.translate(op.bool_expr)
     true_expr = translator.translate(op.true_expr)
@@ -234,6 +246,7 @@ def _floor_divide(translator: ExprTranslator, op: ops.Node) -> str:
         ops.ExtractSecond: _extract_field("second"),  # equivalent to SECOND(timestamp)
         # Other operations
         ops.Literal: _literal,
+        ops.TryCast: _try_cast,
         ops.IfElse: _filter,
         ops.TimestampDiff: _timestamp_diff,
         ops.TimestampFromUNIX: _timestamp_from_unix,

diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py
@@ -1211,14 +1211,16 @@ def test_hash_consistent(backend, alltypes):
             "a",
             "int",
             None,
-            marks=pytest.mark.notyet(["polars"], reason="polars casts to nan"),
+            marks=pytest.mark.notyet(["polars", "flink"], reason="casts to nan"),
         ),
         param(
             datetime.datetime(2023, 1, 1),
             "int",
             None,
             marks=[
-                pytest.mark.never(["clickhouse"], reason="casts to 1672531200"),
+                pytest.mark.never(
+                    ["clickhouse", "flink"], reason="casts to 1672531200"
+                ),
                 pytest.mark.notyet(
                     ["trino"],
                     raises=sa.exc.ProgrammingError,
@@ -1301,14 +1303,14 @@ def test_try_cast_table(con):
             np.isnan,
             marks=[
                 pytest.mark.notyet(
-                    ["clickhouse"], reason="clickhouse casts this to to a number"
+                    ["clickhouse", "polars", "flink"],
+                    reason="casts this to to a number",
                 ),
                 pytest.mark.notyet(
                     ["trino"],
                     raises=sa.exc.ProgrammingError,
                     reason="raises TrinoUserError",
                 ),
-                pytest.mark.notyet(["polars"], reason="polars casts this to a number"),
             ],
         ),
     ],