o 8Di @sddlmZddlZddlmZddlmZddlmZddl m Z ddl m Z m Z mZedZd!d"ddZd#d$ddZd edfd%ddZeddfddZddZdd ZdS)&) annotationsN)exp) OptimizeError)while_changing)find_all_in_scope)flattenrewrite_between uniq_sortsqlglotF expressionexp.Expressiondnfbool max_distanceintc s t|jdddD]x}t|tjrt|drq ||u}|}|jtddt |d}|krBt d|d |Sz| t |fd d}Wn,ty}}z t || ||ro|WYd }~S|WYd }~Sd }~ww|r|}q |S) a Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form. Example: >>> import sqlglot >>> expression = sqlglot.parse_one("(x AND y) OR z") >>> normalize(expression, dnf=False).sql() '(x OR z) AND (y OR z)' Args: expression: expression to normalize dnf: rewrite in disjunctive normal form instead. max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion Returns: sqlglot.Expression: normalized expression cSs t|tjSN) isinstancer ConnectorerT/var/www/Datamplify/venv/lib/python3.10/site-packages/sqlglot/optimizer/normalize.py znormalize..)prunerFcopyrmax_z(Skipping normalization because distance exceeds max c t|Srdistributive_lawrrrrrr1rN)tuplewalkrrr normalizedr transformrnormalization_distanceloggerinforeplacerr)r rrnoderootoriginaldistancerrr%r normalizes:    r2returncs>|rtjtjfntjtjf\}tfddt||D S)av Checks whether a given expression is in a normal form of interest. Example: >>> from sqlglot import parse_one >>> normalized(parse_one("(a AND b) OR c OR (d AND e)"), dnf=True) True >>> normalized(parse_one("(a OR b) AND c")) # Checks CNF by default True >>> normalized(parse_one("a AND (b OR c)"), dnf=True) False Args: expression: The expression to check if it's normalized. dnf: Whether to check if the expression is in Disjunctive Normal Form (DNF). Default: False, i.e. we check if it's in Conjunctive Normal Form (CNF). c3s|]}|VqdSr) find_ancestor).0 connectorancestorrr Ss  znormalized..)rAndOranyr)r rr/rr7rr(@s   r(infr floatcCsNtdd|tjDd }t|||D]}||7}||kr$|Sq|S)a The difference in the number of predicates between a given expression and its normalized form. This is used as an estimate of the cost of the conversion which is exponential in complexity. Example: >>> import sqlglot >>> expression = sqlglot.parse_one("(a AND b) OR (c AND d)") >>> normalization_distance(expression) 4 Args: expression: The expression to compute the normalization distance for. dnf: Whether to check if the expression is in Disjunctive Normal Form (DNF). Default: False, i.e. we check if it's in Conjunctive Normal Form (CNF). max_: stop early if count exceeds this. Returns: The normalization distance. css|]}dVqdS)Nr)r5_rrrr9osz)normalization_distance..r?)sumfind_allrr_predicate_lengths)r rr totallengthrrrr*Xs r*ccs||kr |VdS|}t|tjsdVdS|d7}|j\}}t||r+tjntjrJt||||D]}t||||D]}||Vq?q6dSt||||EdHt||||EdHdS)z Returns a list of predicate lengths when expanded to normalized form. (A AND B) OR C -> [2, 2] because len(A OR C), len(B OR C). Nr?) unnestrrrargsvaluesr:r;rC)r rr depthleftrightabrrrrCys$  rCc s>t|dr|St|d}|krtd|dt|fddr0tjtjfntjtjf\}}t||r|\}}|tjkrKtj ntj }|tjkrVtj ntj } t||rt||rt t | tjt t | tjkr~t|||| St|||| St||rt|||| St||rt|||| S|S)z x OR (y AND z) -> (x OR y) AND (x OR z) (x AND y) OR (y AND z) -> (x OR y) AND (x OR z) AND (y OR y) AND (y OR z) rrzNormalization distance r!cr"rr#rr%rrrrz"distributive_law..)r(r*rrreplace_childrenr;r:runnest_operandsand_or_lenr&rBr _distribute) r rrr1to_expfrom_exprLrM from_functo_funcrr%rr$s(    (  r$csZt|tjrt|fdd|Stt|jtt|jdd}|S)Ncs.tt|jtt|jddS)NFr)r rrJrK)crMrVrWrrrs z_distribute..Fr)rrrrNr rrJrK)rLrMrVrWrrYrrSs rS)Fr )r r rrrr)F)r r rrr3r)r r rrr r>r3r) __future__rloggingr rsqlglot.errorsrsqlglot.helperrsqlglot.optimizer.scopersqlglot.optimizer.simplifyrrr getLoggerr+r2r(r>r*rCr$rSrrrrs        2 ! "