feat(snowflake): implement StructField and struct literals

ibis-project · Jan 19, 2023 · 286a5c3 · 286a5c3
1 parent 314637d
commit 286a5c3
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 11 deletions.
diff --git a/ibis/backends/snowflake/registry.py b/ibis/backends/snowflake/registry.py
@@ -49,7 +49,7 @@ def _literal(t, op):
         return sa.func.date_from_parts(value.year, value.month, value.day)
     elif dtype.is_array():
         return sa.func.array_construct(*value)
-    elif dtype.is_map():
+    elif dtype.is_map() or dtype.is_struct():
         return sa.func.object_construct_keep_null(
             *itertools.chain.from_iterable(value.items())
         )
@@ -118,7 +118,7 @@ def _map(_, op):
     ):
         raise TypeError("Both keys and values of an `ibis.map` call must be literals")
 
-    return sa.object_construct_keep_null(
+    return sa.func.object_construct_keep_null(
         *itertools.chain.from_iterable(zip(keys.value, values.value))
     )
 
@@ -132,7 +132,6 @@ def _map(_, op):
 operation_registry.update(
     {
         ops.JSONGetItem: fixed_arity(sa.func.get, 2),
-        ops.StructField: fixed_arity(sa.func.get, 2),
         ops.StringFind: _string_find,
         ops.MapKeys: unary(sa.func.object_keys),
         ops.MapGet: fixed_arity(
@@ -238,6 +237,10 @@ def _map(_, op):
         ops.TimestampFromUNIX: lambda t, op: sa.func.to_timestamp(
             t.translate(op.arg), _TIMESTAMP_UNITS_TO_SCALE[op.unit]
         ),
+        ops.StructField: lambda t, op: sa.cast(
+            sa.func.parse_json(sa.func.get(t.translate(op.arg), op.field)),
+            t.get_sqla_type(op.output_dtype),
+        ),
     }
 )
 
@@ -255,8 +258,6 @@ def _map(_, op):
     ops.MultiQuantile,
     # ibis.expr.operations.strings
     ops.FindInSet,
-    # ibis.expr.operations.structs
-    ops.StructField,
     # ibis.expr.operations.temporal
     ops.IntervalFromInteger,
     ops.TimestampDiff,

diff --git a/ibis/backends/tests/test_param.py b/ibis/backends/tests/test_param.py
@@ -69,9 +69,7 @@ def test_scalar_param_array(con):
     assert result == len(value)
 
 
-@pytest.mark.notimpl(
-    ["clickhouse", "datafusion", "impala", "postgres", "pyspark", "snowflake"]
-)
+@pytest.mark.notimpl(["clickhouse", "datafusion", "impala", "postgres", "pyspark"])
 @pytest.mark.never(
     ["mysql", "sqlite", "mssql"],
     reason="mysql and sqlite will never implement struct types",

diff --git a/ibis/backends/tests/test_sql.py b/ibis/backends/tests/test_sql.py
@@ -59,7 +59,7 @@ def test_table(con):
     reason="structs not supported in the backend",
 )
 no_struct_literals = mark.notimpl(
-    ["bigquery", "postgres", "snowflake", "mssql"],
+    ["bigquery", "postgres", "mssql"],
     reason="struct literals are not yet implemented",
 )
 not_sql = mark.never(

diff --git a/ibis/backends/tests/test_struct.py b/ibis/backends/tests/test_struct.py
@@ -51,7 +51,7 @@ def test_all_fields(struct, struct_df):
 _NULL_STRUCT_LITERAL = ibis.NA.cast("struct<a: int64, b: string, c: float64>")
 
 
-@pytest.mark.notimpl(["postgres", "snowflake", "bigquery"])
+@pytest.mark.notimpl(["postgres", "bigquery"])
 @pytest.mark.parametrize("field", ["a", "b", "c"])
 def test_literal(con, field):
     query = _STRUCT_LITERAL[field]
@@ -62,7 +62,7 @@ def test_literal(con, field):
     tm.assert_series_equal(result, expected.astype(dtype))
 
 
-@pytest.mark.notimpl(["postgres", "snowflake"])
+@pytest.mark.notimpl(["postgres"])
 @pytest.mark.parametrize("field", ["a", "b", "c"])
 @pytest.mark.notyet(
     ["clickhouse"], reason="clickhouse doesn't support nullable nested types"