feat(snowflake): add more map operations

Author: cpcloud

ci/schema/snowflake.sql

@@ -94,6 +94,13 @@ INSERT INTO array_types ("x", "y", "z", "grouper", "scalar_column", "multi_dim")
     SELECT [2, NULL, 3], ['b', NULL, 'c'], NULL, 'b', 5.0, NULL UNION
     SELECT [4, NULL, NULL, 5], ['d', NULL, NULL, 'e'], [4.0, NULL, NULL, 5.0], 'c', 6.0, [[1, 2, 3]];
 
+CREATE OR REPLACE TABLE map ("kv" OBJECT);
+
+INSERT INTO map ("kv")
+    SELECT object_construct('a', 1, 'b', 2, 'c', 3) UNION
+    SELECT object_construct('d', 4, 'e', 5, 'c', 6);
+
+
 CREATE OR REPLACE TABLE struct ("abc" OBJECT);
 
 INSERT INTO struct ("abc")

ibis/backends/snowflake/registry.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import itertools
+
 import numpy as np
 import sqlalchemy as sa
 from snowflake.sqlalchemy.custom_types import VARIANT
@@ -48,6 +50,10 @@ 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():
+        return sa.func.object_construct_keep_null(
+            *zip(itertools.chain.from_iterable(value.items()))
+        )
     return _postgres_literal(t, op)
 
 
@@ -116,6 +122,19 @@ def _array_slice(t, op):
         ops.StructField: fixed_arity(sa.func.get, 2),
         ops.StringFind: _string_find,
         ops.MapKeys: unary(sa.func.object_keys),
+        ops.MapGet: fixed_arity(
+            lambda arg, key, default: sa.func.coalesce(
+                sa.func.get(arg, key), sa.cast(default, VARIANT)
+            ),
+            3,
+        ),
+        ops.MapContains: fixed_arity(
+            lambda arg, key: sa.func.array_contains(
+                sa.func.cast(key, VARIANT), sa.func.object_keys(arg)
+            ),
+            2,
+        ),
+        ops.MapLength: unary(lambda arg: sa.func.array_size(sa.func.object_keys(arg))),
         ops.BitwiseLeftShift: fixed_arity(sa.func.bitshiftleft, 2),
         ops.BitwiseRightShift: fixed_arity(sa.func.bitshiftright, 2),
         ops.Ln: unary(sa.func.ln),
@@ -196,8 +215,6 @@ def _array_slice(t, op):
     # ibis.expr.operations.array
     ops.ArrayRepeat,
     ops.Unnest,
-    # ibis.expr.operations.maps
-    ops.MapKeys,
     # ibis.expr.operations.reductions
     ops.All,
     ops.Any,

ibis/backends/tests/test_map.py

@@ -5,21 +5,15 @@
 import ibis.expr.datatypes as dt
 
 pytestmark = [
-    pytest.mark.never(["sqlite", "mysql"], reason="No map support"),
+    pytest.mark.never(
+        ["sqlite", "mysql", "mssql", "postgres"], reason="No map support"
+    ),
+    pytest.mark.notyet(
+        ["bigquery", "impala"], reason="backend doesn't implement map types"
+    ),
     pytest.mark.notimpl(
-        [
-            "duckdb",
-            "postgres",
-            "impala",
-            "datafusion",
-            "pyspark",
-            "snowflake",
-            "polars",
-            "mssql",
-        ],
-        reason="Not implemented yet",
+        ["duckdb", "datafusion", "pyspark", "polars"], reason="Not implemented yet"
     ),
-    pytest.mark.notyet(["bigquery"], reason="BigQuery doesn't implement map types"),
 ]
 
 
@@ -49,6 +43,7 @@ def test_literal_map_values(con):
 
 
 @pytest.mark.notimpl(["trino"])
+@pytest.mark.notyet(["snowflake"])
 def test_scalar_isin_literal_map_keys(con):
     mapping = ibis.literal({'a': 1, 'b': 2})
     a = ibis.literal('a')
@@ -78,6 +73,7 @@ def test_map_scalar_contains_key_column(backend, alltypes, df):
     backend.assert_series_equal(result, expected)
 
 
+@pytest.mark.notyet(["snowflake"])
 def test_map_column_contains_key_scalar(backend, alltypes, df):
     expr = ibis.map(ibis.array([alltypes.string_col]), ibis.array([alltypes.int_col]))
     series = df.apply(lambda row: {row['string_col']: row['int_col']}, axis=1)
@@ -88,12 +84,14 @@ def test_map_column_contains_key_scalar(backend, alltypes, df):
     backend.assert_series_equal(result, series)
 
 
+@pytest.mark.notyet(["snowflake"])
 def test_map_column_contains_key_column(backend, alltypes, df):
     expr = ibis.map(ibis.array([alltypes.string_col]), ibis.array([alltypes.int_col]))
     result = expr.contains(alltypes.string_col).name('tmp').execute()
     assert result.all()
 
 
+@pytest.mark.notyet(["snowflake"])
 def test_literal_map_merge(con):
     a = ibis.literal({'a': 0, 'b': 2})
     b = ibis.literal({'a': 1, 'c': 3})
@@ -126,6 +124,7 @@ def test_literal_map_get_broadcast(backend, alltypes, df):
     backend.assert_series_equal(result, expected)
 
 
+@pytest.mark.notyet(["snowflake"])
 def test_map_construction(con, alltypes, df):
     expr = ibis.map(['a', 'b'], [1, 2])
     result = con.execute(expr.name('tmp'))