Skip to content

Commit

Permalink
fix(snowflake): allow loose casting of objects and arrays
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Feb 6, 2023
1 parent 34f3898 commit 1cf8df0
Showing 1 changed file with 25 additions and 1 deletion.
26 changes: 25 additions & 1 deletion ibis/backends/snowflake/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@

import numpy as np
import sqlalchemy as sa
from snowflake.sqlalchemy import ARRAY
from snowflake.sqlalchemy import ARRAY, OBJECT, VARIANT
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import sqltypes
from sqlalchemy.sql.elements import Cast
from sqlalchemy.sql.functions import GenericFunction

import ibis.expr.operations as ops
Expand Down Expand Up @@ -136,6 +137,29 @@ def _arbitrary(t, op):
return t._reduction(sa.func.min, op)


@compiles(Cast, "snowflake")
def compiles_cast(element, compiler, **kw):
arg = compiler.process(element.clause, **kw)
typ = compiler.process(element.type)
if typ in ("OBJECT", "ARRAY"):
return f"IFF(IS_{typ}({arg}), {arg}, NULL)"
return compiler.visit_cast(element, **kw)


@compiles(sa.Text, "snowflake")
@compiles(sa.TEXT, "snowflake")
@compiles(sa.VARCHAR, "snowflake")
def compiles_string(element, compiler, **kw):
return "VARCHAR"


@compiles(OBJECT, "snowflake")
@compiles(ARRAY, "snowflake")
@compiles(VARIANT, "snowflake")
def compiles_object_type(element, compiler, **kw):
return type(element).__name__.upper()


class _flatten(GenericFunction):
def __init__(self, arg, *, type: sa.types.TypeEngine) -> None:
super().__init__(arg)
Expand Down

0 comments on commit 1cf8df0

Please sign in to comment.