Cannot annotate type of unnest for postgres

[eric-zhu]

Version

$ pip list | grep sqlglot
sqlglot               21.0.1

Fully reproducible code snippet

import sqlglot
from sqlglot.optimizer.annotate_types import annotate_types
from sqlglot.optimizer.qualify import qualify

schema = {
    "customers": {
        "access_errors": "text[]",
    }
}
sql_query = "select unnest(access_errors) as errors from customers"
expression = sqlglot.parse_one(sql_query, dialect="postgres")
qualified_expr = qualify(expression, dialect="postgres", schema=schema)
annotated_expr = annotate_types(qualified_expr, schema=schema)
ex = annotated_expr.expressions[0]
print(f"column '{ex.alias_or_name}' type: {ex.type}")

Output: column 'errors' type: UNKNOWN.
Expected: column 'errors' type: TEXT

import sqlglot
from sqlglot.optimizer.annotate_types import annotate_types
from sqlglot.optimizer.qualify import qualify

sql_query = "select unnest('{A,B,C,D}'::text[]) category"
expression = sqlglot.parse_one(sql_query, dialect="postgres")
qualified_expr = qualify(expression, dialect="postgres")
annotated_expr = annotate_types(qualified_expr)
ex = annotated_expr.expressions[0]
print(f"column '{ex.alias_or_name}' type: {ex.type}")

Output: column 'category' type: UNKNOWN
Expected: column 'category' type: TEXT

[tobymao]

unnest is currently unsupported but your code snippet has an error, you need to pass schema to annotate types

[eric-zhu]

unnest is currently unsupported but your code snippet has an error, you need to pass schema to annotate types

Thanks for pointing out the error. I've update the code snippet. Output is the same.

Is it easy to add support for unnest? or do you have a plan to add it? Thanks.

[tobymao]

it’s easy, but i ran into an obscure bug that i need to fix first before i can get it in

[eric-zhu]

Thanks a lot for the quick fix.

[tobymao]

you're welcome!