Best way to split a dynamically long "where in" predicate to multiple predicates

[shaiv-syg]

We encountered an issue when we build our SQL statement dynamically with a very long list.

When we generate a statement like this:

SELECT *
FROM X
WHERE Y IN  <very_long_list>

We receive an exception Error: maximum number of expressions in a list exceeded, expected at most 16,384

We thought that we can try to split the IN predicate to multiple predicates with an OR something like this:

SELECT *
FROM X
WHERE Y IN  <partial_list_a> OR Y IN <partial_list_b>

I thought of using the transformer pattern but I can't seem to figure out how to implement the body of the transform function.
Can anyone point me to the right direction?

from sqlglot import exp, parse_one

very_long_list = (1, 2, 3, 4, 5)
query = f"""
SELECT *
FROM X
WHERE Y IN {very_long_list}
"""

def transformer(node):
    if isinstance(node, exp.In):
        # what to do here...
        pass
    return node

parsed = parse_one(query)
transformed_tree = parsed.transform(transformer)
transformed_tree.sql()

[tobymao]

where are you getting the max number of expressions error? is this from your database?

the transform should work, just return an or statement,

exp.or_(…., …)

[shaiv-syg]

Yes I'm receiving the error from Snowflake there is a limit of 16,384 items in any type of list.
What I don't understand is if I split the expressions list by chunks how do I create a new node?
I tried the following without success:

from sqlglot import exp, parse_one

MAX_LIST_COUNT = 2 # 16384
def transformer(node):
    if isinstance(node, exp.In):
        list_items = [node.expressions[i:i + MAX_LIST_COUNT] for i in range(0, len(node.expressions), MAX_LIST_COUNT)]
        # [[(LITERAL this: 1, is_string: False), (LITERAL this: 2, is_string: False)], [(LITERAL this: 3, is_string: False), (LITERAL this: 4, is_string: False)], [(LITERAL this: 5, is_string: False)]
        # how can I create a new In stamtment with this expressions list?
    return node


very_long_list = (1, 2, 3, 4, 5)

query = f"""
SELECT *
FROM X
WHERE Y IN {very_long_list}
"""
parsed = parse_one(query)
transformed_tree = parsed.transform(transformer)
transformed_tree.sql()

[tobymao]

by passing it into exp.or_

[shaiv-syg]

Got it to work!
Thanks tobymao - your package is a so neat!

from sqlglot import exp, parse_one

MAX_LIST_COUNT = 2 # 16384
def transformer(node):
    if isinstance(node, exp.In):
        list_items = [node.expressions[i:i + MAX_LIST_COUNT] for i in range(0, len(node.expressions), MAX_LIST_COUNT)]
        splitted = [exp.In(this=exp.Column(this=node.find(exp.Column)), expressions=items) for items in list_items]
        expression = splitted[0]
        return expression.or_(*splitted[1:])
    return node


very_long_list = (1, 2, 3, 4, 5)

query = f"""
SELECT *
FROM X
WHERE ZZZ IN {very_long_list}
"""
parsed = parse_one(query)
transformed_tree = parsed.transform(transformer)
transformed_tree.sql()

Generate the following

'SELECT * FROM X WHERE ZZZ IN (1, 2) OR ZZZ IN (3, 4) OR ZZZ IN (5)'