diff --git a/ibis/backends/snowflake/compiler.py b/ibis/backends/snowflake/compiler.py index ff80e20e5116..33d4ad325e2e 100644 --- a/ibis/backends/snowflake/compiler.py +++ b/ibis/backends/snowflake/compiler.py @@ -48,8 +48,6 @@ class SnowflakeCompiler(SQLGlotCompiler): } UNSUPPORTED_OPS = ( - ops.ArrayMap, - ops.ArrayFilter, ops.RowID, ops.MultiQuantile, ops.IntervalFromInteger, @@ -639,3 +637,22 @@ def visit_Sample( seed=None if seed is None else sge.convert(seed), ) return sg.select(STAR).from_(sample) + + def visit_ArrayMap(self, op, *, arg, param, body): + return self.f.transform(arg, sge.Lambda(this=body, expressions=[param])) + + def visit_ArrayFilter(self, op, *, arg, param, body): + return self.f.filter( + arg, + sge.Lambda( + this=sg.and_( + body, + # necessary otherwise null values are treated as JSON nulls + # instead of SQL NULLs + self.cast(sg.to_identifier(param), op.dtype.value_type).is_( + sg.not_(NULL) + ), + ), + expressions=[param], + ), + ) diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py index 7340540f7269..d78b9bf326e1 100644 --- a/ibis/backends/tests/test_array.py +++ b/ibis/backends/tests/test_array.py @@ -418,14 +418,7 @@ def test_array_slice(backend, start, stop): @builtin_array @pytest.mark.notimpl( - [ - "datafusion", - "flink", - "polars", - "snowflake", - "sqlite", - ], - raises=com.OperationNotDefinedError, + ["datafusion", "flink", "polars", "sqlite"], raises=com.OperationNotDefinedError ) @pytest.mark.broken( ["risingwave"], @@ -465,6 +458,7 @@ def test_array_slice(backend, start, stop): functools.partial(lambda x, y: x + y, y=1), ibis._ + 1, ], + ids=["lambda", "partial", "deferred"], ) @pytest.mark.broken( ["risingwave"], @@ -485,14 +479,7 @@ def test_array_map(con, input, output, func): @builtin_array @pytest.mark.notimpl( - [ - "dask", - "datafusion", - "flink", - "pandas", - "polars", - "snowflake", - ], + ["dask", "datafusion", "flink", "pandas", "polars"], raises=com.OperationNotDefinedError, ) @pytest.mark.notimpl( @@ -533,6 +520,7 @@ def test_array_map(con, input, output, func): functools.partial(lambda x, y: x > y, y=1), ibis._ > 1, ], + ids=["lambda", "partial", "deferred"], ) def test_array_filter(con, input, output, predicate): t = ibis.memtable(input, schema=ibis.schema(dict(a="!array"))) @@ -1138,14 +1126,7 @@ def test_unnest_empty_array(con): @builtin_array @pytest.mark.notimpl( - [ - "datafusion", - "flink", - "polars", - "snowflake", - "dask", - "pandas", - ], + ["datafusion", "flink", "polars", "dask", "pandas"], raises=com.OperationNotDefinedError, ) @pytest.mark.notimpl(["sqlite"], raises=com.UnsupportedBackendType) @@ -1166,16 +1147,7 @@ def test_array_map_with_conflicting_names(backend, con): @builtin_array @pytest.mark.notimpl( - [ - "datafusion", - "flink", - "polars", - "snowflake", - "sqlite", - "dask", - "pandas", - "sqlite", - ], + ["datafusion", "flink", "polars", "sqlite", "dask", "pandas", "sqlite"], raises=com.OperationNotDefinedError, ) def test_complex_array_map(con):