How to check if logical expression allows not more than N partitions

[evyasonov]

Hi

I want to check if a user scans not too much partitions. Partitions are represented by dates. Other words, I need to scan WHERE part and check that length if the range of dates is not too big. For instanse, for WHERE date BETWEEN '2023-01-01' AND '2023-01-02' the answer should be 2.

But the problem is that WHERE part might be complex, such as:

WHERE a = 10 AND (date BETWEEN '2023-01-01' AND '2023-01-02' OR user_id IN (SELECT id FROM users))

In this scenario the answer is 'infinity'. So I can not just find date and get the limits.

And I don't see an obvious way to check the length of the range of dates. So I'm asking for help.

Any ideas?

[georgesittas]

One way to tackle this is to call find_all on the where argument of your expression, looking for exp.Between. Then, you can search for only those filters you care about, i.e. those involving date, and manually checking the length of the date range by using something like datetime.

[evyasonov]

Mmmm... I thought that you suggest to transform logical impression into disjunctive normal form and then extract partition columns from each simple conjunction and union it. And I hoped that you write "we have optimizer, you can use a function from it to transform into dnf".

Do you know if optimizer has a function to transform to dnf?

Update: found optimizer.normalize(dnf=True)
Update:

import sqlglot, sqlglot.optimizer
e = sqlglot.parse_one("a = 10 AND (date BETWEEN '2023-01-01' AND '2023-01-02' OR user_id IN (SELECT id FROM users))")
sqlglot.optimizer.normalize.normalize(e, dnf=True).sql()

# (a = 10 AND date <= '2023-01-02' AND date >= '2023-01-01') OR (a = 10 AND user_id IN (SELECT id FROM users))

and then extract dates from each ANDs and do all the magic

[georgesittas]

Maybe I understood the initial intent incorrectly, but you're right, you can convert it to DNF and check whether you can "bypass" the BETWEEN filter, i.e. by having another, independent filter, connected via disjunction. Having this information will only enable you to detect whether it's possible for the range to be "infinity", though. Is that enough for your use case?

[evyasonov]

Yes...
In case one of ANDs does not contain partition column, then the result is infinity.
In case all ANDs contains partition columns, then I extract it as intervals and see a size of union.

Something like this:

import sqlglot, sqlglot.optimizer
e = sqlglot.parse_one("a = 10 AND (date BETWEEN '2023-01-01' AND '2023-01-02' OR user_id IN (SELECT id FROM users))")
sqlglot.optimizer.normalize.normalize(e, dnf=True).sql()

ands = []
while type(e) == exp.Or:
    ands.append(e.expression)
    e = e.this
ands.append(e)

for e in ands:
    # Leave >, <, >=, <=, =, IN only. Replace everything else to some constant. Make sure there's no subqueries, functions, lamdas, selects.
    # if no partition column in `e`, return infinity
    # extact expressions which contains partition column as an interval
    # calculate length of the union

Can you please advice what is the best way to clean all except >, <, >=, <=, =, IN ?

[georgesittas]

You can probably use transform for this.

[evyasonov]

done something like this (haven't been tested, but might be helpful to someone in future):

import sqlglot, sqlglot.optimizer

e = sqlglot.parse_one("a = 10 AND (date BETWEEN '2023-01-01' AND '2023-02-02' OR user_id IN (SELECT id FROM users) OR b = 20) AND c = 30 AND date BETWEEN '2022-01-01' AND '2023-01-14'")
e = sqlglot.optimizer.normalize.normalize(e, dnf=True)

ands = []
while type(e) == exp.Or:
    ands.append(e.expression)
    e = e.this
ands.append(e)

import portion

partition_column = 'date'

result_interval = portion.empty() 
for e in ands:
    and_parts = []
    while type(e) in {exp.And, exp.Paren}:
        if type(e) == exp.And:
            and_parts.append(e.expression)
        e = e.this
    and_parts.append(e)
    
    interval = portion.open(-portion.inf, portion.inf)
    for p in and_parts:
        if type(p) in {exp.EQ, exp.GT, exp.GTE, exp.LT, exp.LTE}:
            if type(p.this) == exp.Column and type(p.expression) in {exp.Literal, exp.Cast}:
                column = p.this
                value  = p.expression
            elif type(p.this) in {exp.Literal, exp.Cast} and type(p.expression) == exp.Column:
                column = p.expression
                value  = p.this
            else:
                continue
        elif type(p) == exp.In:
            column = p.this
            value  = p.expressions
        else:
            continue

        if type(column.this) == exp.Identifier and column.this.this == partition_column:
            if type(value) == exp.Cast and type(value.this) == exp.Literal:
                value = value.this
            
            if type(value) == exp.Literal:
                value = value.this
                if type(p) == exp.EQ:
                    interval &= portion.singleton(value)
                elif type(p) == exp.GT:
                    interval &= portion.open(value, portion.inf)
                elif type(p) == exp.GTE:
                    interval &= portion.closedopen(value, portion.inf)
                elif type(p) == exp.LT:
                    interval &= portion.open(-portion.inf, value)
                elif type(p) == exp.LTE:
                    interval &= portion.openclosed(-portion.inf, value)
            elif type(value) == list:
                i = portion.empty()
                for v in value:
                    if type(v) == exp.Cast and type(v.this) == exp.Literal:
                        v = v.this
                    if type(v) == exp.Literal:
                        i |= portion.singleton(v.this)
                if i:
                    interval &= i
            
    result_interval |= interval

result_interval

@georgesittas can you please comment what can I improve here?