1
0
Fork 0

Adding upstream version 18.13.0.

Signed-off-by: Daniel Baumann <daniel@debian.org>
This commit is contained in:
Daniel Baumann 2025-02-13 21:07:20 +01:00
parent e833f2baa5
commit fc6bad5705
Signed by: daniel
GPG key ID: FBB4F0E80A80222F
76 changed files with 21248 additions and 19605 deletions

View file

@ -1,6 +1,36 @@
Changelog
=========
## [v18.12.0] - 2023-10-10
### :boom: BREAKING CHANGES
- due to [`28308da`](https://github.com/tobymao/sqlglot/commit/28308dae5ab7a0aff3ca2afeff88fa4554babba3) - support spark trycast, treat databricks as strict_cast closes [#2389](https://github.com/tobymao/sqlglot/pull/2389) *(commit by [@tobymao](https://github.com/tobymao))*:
support spark trycast, treat databricks as strict_cast closes #2389
- due to [`c7c3869`](https://github.com/tobymao/sqlglot/commit/c7c3869b01e984a243c071660f27a2c6c4863892) - add explode outer and change hiearchy are explosions closes [#2393](https://github.com/tobymao/sqlglot/pull/2393) *(commit by [@tobymao](https://github.com/tobymao))*:
add explode outer and change hiearchy are explosions closes #2393
### :sparkles: New Features
- [`f4f8366`](https://github.com/tobymao/sqlglot/commit/f4f8366f6f761fefd72cd2a1ee4c462a0e18ec42) - **schema**: add method to check if column exists *(PR [#2381](https://github.com/tobymao/sqlglot/pull/2381) by [@GeorgeSittas](https://github.com/GeorgeSittas))*
- [`da2c6f1`](https://github.com/tobymao/sqlglot/commit/da2c6f167be2de8b61ab9b0fb60c2bc2b6b24408) - **optimizer**: simplify CONCAT_WS *(PR [#2383](https://github.com/tobymao/sqlglot/pull/2383) by [@GeorgeSittas](https://github.com/GeorgeSittas))*
- [`cca58dd`](https://github.com/tobymao/sqlglot/commit/cca58dd2e7a45d1150b37f8e76baa3571fed8135) - **optimizer**: propagate constants *(PR [#2386](https://github.com/tobymao/sqlglot/pull/2386) by [@GeorgeSittas](https://github.com/GeorgeSittas))*
- [`28308da`](https://github.com/tobymao/sqlglot/commit/28308dae5ab7a0aff3ca2afeff88fa4554babba3) - support spark trycast, treat databricks as strict_cast closes [#2389](https://github.com/tobymao/sqlglot/pull/2389) *(commit by [@tobymao](https://github.com/tobymao))*
- [`c7c3869`](https://github.com/tobymao/sqlglot/commit/c7c3869b01e984a243c071660f27a2c6c4863892) - add explode outer and change hiearchy are explosions closes [#2393](https://github.com/tobymao/sqlglot/pull/2393) *(commit by [@tobymao](https://github.com/tobymao))*
### :bug: Bug Fixes
- [`0fb1652`](https://github.com/tobymao/sqlglot/commit/0fb1652784845f083b87e952644a8a9790d28d8e) - replace executor None values with np.NaN to silence Pandas>2.1 warnings *(PR [#2384](https://github.com/tobymao/sqlglot/pull/2384) by [@GeorgeSittas](https://github.com/GeorgeSittas))*
- [`a849794`](https://github.com/tobymao/sqlglot/commit/a8497944424a524fdc1c9fdb9e10aa2f3558bdd5) - **mysql**: move parsing logic for JSON_TABLE to base parser *(PR [#2387](https://github.com/tobymao/sqlglot/pull/2387) by [@GeorgeSittas](https://github.com/GeorgeSittas))*
- [`e08c1c0`](https://github.com/tobymao/sqlglot/commit/e08c1c0e8a5a2fbc46872d163eac52071b49a1cf) - correctly handle agg subexpressions with no selections *(PR [#2390](https://github.com/tobymao/sqlglot/pull/2390) by [@ginter](https://github.com/ginter))*
- [`8afa7a1`](https://github.com/tobymao/sqlglot/commit/8afa7a12ca35400133a49571b49d001bdb526188) - **schema**: don't specialize type variable in MappingSchema *(PR [#2394](https://github.com/tobymao/sqlglot/pull/2394) by [@GeorgeSittas](https://github.com/GeorgeSittas))*
## [v18.11.6] - 2023-10-06
### :bug: Bug Fixes
- [`ef12aa7`](https://github.com/tobymao/sqlglot/commit/ef12aa7b4c24b431dbedcf917f61d18a89dc3a0f) - normalize_identifiers parses identifier strings *(commit by [@tobymao](https://github.com/tobymao))*
## [v18.11.5] - 2023-10-06
### :sparkles: New Features
- [`ae27e46`](https://github.com/tobymao/sqlglot/commit/ae27e46cf60bbbcb456997afe942a6e8ab9d03c1) - **spark**: from_utc_timestamp -> at time zone *(commit by [@tobymao](https://github.com/tobymao))*
@ -1668,3 +1698,5 @@ Changelog
[v18.11.3]: https://github.com/tobymao/sqlglot/compare/v18.11.2...v18.11.3
[v18.11.4]: https://github.com/tobymao/sqlglot/compare/v18.11.3...v18.11.4
[v18.11.5]: https://github.com/tobymao/sqlglot/compare/v18.11.4...v18.11.5
[v18.11.6]: https://github.com/tobymao/sqlglot/compare/v18.11.5...v18.11.6
[v18.12.0]: https://github.com/tobymao/sqlglot/compare/v18.11.6...v18.12.0

File diff suppressed because one or more lines are too long

View file

@ -76,8 +76,8 @@
</span><span id="L-12"><a href="#L-12"><span class="linenos">12</span></a><span class="n">__version_tuple__</span><span class="p">:</span> <span class="n">VERSION_TUPLE</span>
</span><span id="L-13"><a href="#L-13"><span class="linenos">13</span></a><span class="n">version_tuple</span><span class="p">:</span> <span class="n">VERSION_TUPLE</span>
</span><span id="L-14"><a href="#L-14"><span class="linenos">14</span></a>
</span><span id="L-15"><a href="#L-15"><span class="linenos">15</span></a><span class="n">__version__</span> <span class="o">=</span> <span class="n">version</span> <span class="o">=</span> <span class="s1">&#39;18.11.5&#39;</span>
</span><span id="L-16"><a href="#L-16"><span class="linenos">16</span></a><span class="n">__version_tuple__</span> <span class="o">=</span> <span class="n">version_tuple</span> <span class="o">=</span> <span class="p">(</span><span class="mi">18</span><span class="p">,</span> <span class="mi">11</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span>
</span><span id="L-15"><a href="#L-15"><span class="linenos">15</span></a><span class="n">__version__</span> <span class="o">=</span> <span class="n">version</span> <span class="o">=</span> <span class="s1">&#39;18.12.0&#39;</span>
</span><span id="L-16"><a href="#L-16"><span class="linenos">16</span></a><span class="n">__version_tuple__</span> <span class="o">=</span> <span class="n">version_tuple</span> <span class="o">=</span> <span class="p">(</span><span class="mi">18</span><span class="p">,</span> <span class="mi">12</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
</span></pre></div>
@ -97,7 +97,7 @@
<section id="version">
<div class="attr variable">
<span class="name">version</span><span class="annotation">: str</span> =
<span class="default_value">&#39;18.11.5&#39;</span>
<span class="default_value">&#39;18.12.0&#39;</span>
</div>
@ -109,7 +109,7 @@
<section id="version_tuple">
<div class="attr variable">
<span class="name">version_tuple</span><span class="annotation">: object</span> =
<span class="default_value">(18, 11, 5)</span>
<span class="default_value">(18, 12, 0)</span>
</div>

View file

@ -777,7 +777,7 @@
<div class="attr function">
<span class="def">def</span>
<span class="name">createDataFrame</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">data</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712877862496&#39;</span><span class="o">&gt;</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712877862496&#39;</span><span class="o">&gt;</span><span class="p">],</span> <span class="n">Tuple</span><span class="p">]]</span>,</span><span class="param"> <span class="n">schema</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712877785152&#39;</span><span class="o">&gt;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">samplingRatio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">verifySchema</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<span class="name">createDataFrame</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">data</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622968000&#39;</span><span class="o">&gt;</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622968000&#39;</span><span class="o">&gt;</span><span class="p">],</span> <span class="n">Tuple</span><span class="p">]]</span>,</span><span class="param"> <span class="n">schema</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622557920&#39;</span><span class="o">&gt;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">samplingRatio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">verifySchema</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<label class="view-source-button" for="SparkSession.createDataFrame-view-source"><span>View Source</span></label>
@ -1829,7 +1829,7 @@
<input id="DataFrame.__init__-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
<div class="attr function">
<span class="name">DataFrame</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="n">spark</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712880355936&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">expression</span><span class="p">:</span> <span class="n"><a href="../expressions.html#Select">sqlglot.expressions.Select</a></span>,</span><span class="param"> <span class="n">branch_id</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">sequence_id</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">last_op</span><span class="p">:</span> <span class="n">sqlglot</span><span class="o">.</span><span class="n">dataframe</span><span class="o">.</span><span class="n">sql</span><span class="o">.</span><span class="n">operations</span><span class="o">.</span><span class="n">Operation</span> <span class="o">=</span> <span class="o">&lt;</span><span class="n">Operation</span><span class="o">.</span><span class="n">INIT</span><span class="p">:</span> <span class="o">-</span><span class="mi">1</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">pending_hints</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">output_expression_container</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712880739504&#39;</span><span class="o">&gt;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span>)</span>
<span class="name">DataFrame</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="n">spark</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180625818192&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">expression</span><span class="p">:</span> <span class="n"><a href="../expressions.html#Select">sqlglot.expressions.Select</a></span>,</span><span class="param"> <span class="n">branch_id</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">sequence_id</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">last_op</span><span class="p">:</span> <span class="n">sqlglot</span><span class="o">.</span><span class="n">dataframe</span><span class="o">.</span><span class="n">sql</span><span class="o">.</span><span class="n">operations</span><span class="o">.</span><span class="n">Operation</span> <span class="o">=</span> <span class="o">&lt;</span><span class="n">Operation</span><span class="o">.</span><span class="n">INIT</span><span class="p">:</span> <span class="o">-</span><span class="mi">1</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">pending_hints</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">output_expression_container</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180626003792&#39;</span><span class="o">&gt;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span>)</span>
<label class="view-source-button" for="DataFrame.__init__-view-source"><span>View Source</span></label>
@ -2018,7 +2018,7 @@
<div class="attr function">
<span class="def">def</span>
<span class="name">sql</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">dialect</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712876492592&#39;</span><span class="o">&gt;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">optimize</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span>:</span></span>
<span class="name">sql</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">dialect</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622657328&#39;</span><span class="o">&gt;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">optimize</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span>:</span></span>
<label class="view-source-button" for="DataFrame.sql-view-source"><span>View Source</span></label>
@ -2773,7 +2773,7 @@ is unlikely to come up.</p>
<div class="decorator">@operation(Operation.FROM)</div>
<span class="def">def</span>
<span class="name">fillna</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">value</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712876223808&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">subset</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">NoneType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<span class="name">fillna</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">value</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621314368&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">subset</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">NoneType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<label class="view-source-button" for="DataFrame.fillna-view-source"><span>View Source</span></label>
@ -2842,7 +2842,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="decorator">@operation(Operation.FROM)</div>
<span class="def">def</span>
<span class="name">replace</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">to_replace</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">bool</span><span class="p">,</span> <span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Dict</span><span class="p">]</span>,</span><span class="param"> <span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">bool</span><span class="p">,</span> <span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">subset</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">Collection</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712876804128&#39;</span><span class="o">&gt;</span><span class="p">],</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712876804128&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<span class="name">replace</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">to_replace</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">bool</span><span class="p">,</span> <span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Dict</span><span class="p">]</span>,</span><span class="param"> <span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">bool</span><span class="p">,</span> <span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>,</span><span class="param"> <span class="n">subset</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">Collection</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621772976&#39;</span><span class="o">&gt;</span><span class="p">],</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621772976&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<label class="view-source-button" for="DataFrame.replace-view-source"><span>View Source</span></label>
@ -3047,7 +3047,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="decorator">@operation(Operation.NO_OP)</div>
<span class="def">def</span>
<span class="name">repartition</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">numPartitions</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875138096&#39;</span><span class="o">&gt;</span><span class="p">]</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712874854672&#39;</span><span class="o">&gt;</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<span class="name">repartition</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">numPartitions</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622057408&#39;</span><span class="o">&gt;</span><span class="p">]</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622201840&#39;</span><span class="o">&gt;</span></span><span class="return-annotation">) -> <span class="n"><a href="#DataFrame">DataFrame</a></span>:</span></span>
<label class="view-source-button" for="DataFrame.repartition-view-source"><span>View Source</span></label>
@ -3765,7 +3765,7 @@ and check if it matches the type of the value provided. If not then make it null
<input id="Column.__init__-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
<div class="attr function">
<span class="name">Column</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="n">expression</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712878583872&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span></span>)</span>
<span class="name">Column</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="n">expression</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180623815392&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span></span>)</span>
<label class="view-source-button" for="Column.__init__-view-source"><span>View Source</span></label>
@ -3809,7 +3809,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="decorator">@classmethod</div>
<span class="def">def</span>
<span class="name">ensure_col</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712874910496&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">ensure_col</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622248832&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">,</span> <span class="n">NoneType</span><span class="p">]</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.ensure_col-view-source"><span>View Source</span></label>
@ -3830,7 +3830,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="decorator">@classmethod</div>
<span class="def">def</span>
<span class="name">ensure_cols</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">args</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875482880&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n">List</span><span class="p">[</span><span class="n"><a href="#Column">Column</a></span><span class="p">]</span>:</span></span>
<span class="name">ensure_cols</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">args</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180622270880&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n">List</span><span class="p">[</span><span class="n"><a href="#Column">Column</a></span><span class="p">]</span>:</span></span>
<label class="view-source-button" for="Column.ensure_cols-view-source"><span>View Source</span></label>
@ -3851,7 +3851,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="decorator">@classmethod</div>
<span class="def">def</span>
<span class="name">invoke_anonymous_function</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">column</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875154096&#39;</span><span class="o">&gt;</span><span class="p">]</span>,</span><span class="param"> <span class="n">func_name</span><span class="p">:</span> <span class="nb">str</span>,</span><span class="param"> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875032880&#39;</span><span class="o">&gt;</span><span class="p">]</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">invoke_anonymous_function</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">column</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620601152&#39;</span><span class="o">&gt;</span><span class="p">]</span>,</span><span class="param"> <span class="n">func_name</span><span class="p">:</span> <span class="nb">str</span>,</span><span class="param"> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620379088&#39;</span><span class="o">&gt;</span><span class="p">]</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.invoke_anonymous_function-view-source"><span>View Source</span></label>
@ -3878,7 +3878,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="decorator">@classmethod</div>
<span class="def">def</span>
<span class="name">invoke_expression_over_column</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">column</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875027344&#39;</span><span class="o">&gt;</span><span class="p">]</span>,</span><span class="param"> <span class="n">callable_expression</span><span class="p">:</span> <span class="n">Callable</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">invoke_expression_over_column</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="n">column</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620187376&#39;</span><span class="o">&gt;</span><span class="p">]</span>,</span><span class="param"> <span class="n">callable_expression</span><span class="p">:</span> <span class="n">Callable</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.invoke_expression_over_column-view-source"><span>View Source</span></label>
@ -3915,7 +3915,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="attr function">
<span class="def">def</span>
<span class="name">binary_op</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">klass</span><span class="p">:</span> <span class="n">Callable</span>,</span><span class="param"> <span class="n">other</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875319184&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">binary_op</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">klass</span><span class="p">:</span> <span class="n">Callable</span>,</span><span class="param"> <span class="n">other</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620223024&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.binary_op-view-source"><span>View Source</span></label>
@ -3936,7 +3936,7 @@ and check if it matches the type of the value provided. If not then make it null
<div class="attr function">
<span class="def">def</span>
<span class="name">inverse_binary_op</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">klass</span><span class="p">:</span> <span class="n">Callable</span>,</span><span class="param"> <span class="n">other</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875330432&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">inverse_binary_op</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">klass</span><span class="p">:</span> <span class="n">Callable</span>,</span><span class="param"> <span class="n">other</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620447264&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="o">**</span><span class="n">kwargs</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.inverse_binary_op-view-source"><span>View Source</span></label>
@ -4502,7 +4502,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="attr function">
<span class="def">def</span>
<span class="name">isin</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875557696&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">Iterable</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875557696&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">):</span></span>
<span class="name">isin</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620677120&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">Iterable</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620677120&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">):</span></span>
<label class="view-source-button" for="Column.isin-view-source"><span>View Source</span></label>
@ -4523,7 +4523,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="attr function">
<span class="def">def</span>
<span class="name">between</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">lowerBound</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875652000&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">upperBound</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712875706544&#39;</span><span class="o">&gt;</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">between</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">lowerBound</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620785216&#39;</span><span class="o">&gt;</span>,</span><span class="param"> <span class="n">upperBound</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620872528&#39;</span><span class="o">&gt;</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.between-view-source"><span>View Source</span></label>
@ -4558,7 +4558,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="attr function">
<span class="def">def</span>
<span class="name">over</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">window</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712876108208&#39;</span><span class="o">&gt;</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<span class="name">over</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="n">window</span><span class="p">:</span> <span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621601872&#39;</span><span class="o">&gt;</span></span><span class="return-annotation">) -> <span class="n"><a href="#Column">Column</a></span>:</span></span>
<label class="view-source-button" for="Column.over-view-source"><span>View Source</span></label>
@ -4803,7 +4803,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="decorator">@classmethod</div>
<span class="def">def</span>
<span class="name">partitionBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712872747136&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712872747136&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<span class="name">partitionBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621091568&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621091568&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<label class="view-source-button" for="Window.partitionBy-view-source"><span>View Source</span></label>
@ -4824,7 +4824,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="decorator">@classmethod</div>
<span class="def">def</span>
<span class="name">orderBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712872888592&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712872888592&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<span class="name">orderBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">cls</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621079648&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180621079648&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<label class="view-source-button" for="Window.orderBy-view-source"><span>View Source</span></label>
@ -5064,7 +5064,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="attr function">
<span class="def">def</span>
<span class="name">partitionBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712873058720&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712873058720&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<span class="name">partitionBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620960832&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180620960832&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<label class="view-source-button" for="WindowSpec.partitionBy-view-source"><span>View Source</span></label>
@ -5091,7 +5091,7 @@ Sqlglot doesn't currently replicate this class so it only accepts a string</p>
<div class="attr function">
<span class="def">def</span>
<span class="name">orderBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712873089328&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140712873089328&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<span class="name">orderBy</span><span class="signature pdoc-code multiline">(<span class="param"> <span class="bp">self</span>,</span><span class="param"> <span class="o">*</span><span class="n">cols</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180619970784&#39;</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="o">&lt;</span><span class="n">MagicMock</span> <span class="nb">id</span><span class="o">=</span><span class="s1">&#39;140180619970784&#39;</span><span class="o">&gt;</span><span class="p">]]</span></span><span class="return-annotation">) -> <span class="n"><a href="#WindowSpec">WindowSpec</a></span>:</span></span>
<label class="view-source-button" for="WindowSpec.orderBy-view-source"><span>View Source</span></label>

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -1032,6 +1032,7 @@ Default: True</li>
<dd id="Trino.Generator.jsonarray_sql" class="function"><a href="../generator.html#Generator.jsonarray_sql">jsonarray_sql</a></dd>
<dd id="Trino.Generator.jsonarrayagg_sql" class="function"><a href="../generator.html#Generator.jsonarrayagg_sql">jsonarrayagg_sql</a></dd>
<dd id="Trino.Generator.jsoncolumndef_sql" class="function"><a href="../generator.html#Generator.jsoncolumndef_sql">jsoncolumndef_sql</a></dd>
<dd id="Trino.Generator.jsonschema_sql" class="function"><a href="../generator.html#Generator.jsonschema_sql">jsonschema_sql</a></dd>
<dd id="Trino.Generator.jsontable_sql" class="function"><a href="../generator.html#Generator.jsontable_sql">jsontable_sql</a></dd>
<dd id="Trino.Generator.openjsoncolumndef_sql" class="function"><a href="../generator.html#Generator.openjsoncolumndef_sql">openjsoncolumndef_sql</a></dd>
<dd id="Trino.Generator.openjson_sql" class="function"><a href="../generator.html#Generator.openjson_sql">openjson_sql</a></dd>

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because it is too large Load diff

View file

@ -577,22 +577,30 @@
</span><span id="L-441"><a href="#L-441"><span class="linenos">441</span></a>
</span><span id="L-442"><a href="#L-442"><span class="linenos">442</span></a>
</span><span id="L-443"><a href="#L-443"><span class="linenos">443</span></a><span class="k">def</span> <span class="nf">merge_ranges</span><span class="p">(</span><span class="n">ranges</span><span class="p">:</span> <span class="n">t</span><span class="o">.</span><span class="n">List</span><span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">A</span><span class="p">]])</span> <span class="o">-&gt;</span> <span class="n">t</span><span class="o">.</span><span class="n">List</span><span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">A</span><span class="p">]]:</span>
</span><span id="L-444"><a href="#L-444"><span class="linenos">444</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">ranges</span><span class="p">:</span>
</span><span id="L-445"><a href="#L-445"><span class="linenos">445</span></a> <span class="k">return</span> <span class="p">[]</span>
</span><span id="L-446"><a href="#L-446"><span class="linenos">446</span></a>
</span><span id="L-447"><a href="#L-447"><span class="linenos">447</span></a> <span class="n">ranges</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">ranges</span><span class="p">)</span>
</span><span id="L-448"><a href="#L-448"><span class="linenos">448</span></a>
</span><span id="L-449"><a href="#L-449"><span class="linenos">449</span></a> <span class="n">merged</span> <span class="o">=</span> <span class="p">[</span><span class="n">ranges</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span>
</span><span id="L-450"><a href="#L-450"><span class="linenos">450</span></a>
</span><span id="L-451"><a href="#L-451"><span class="linenos">451</span></a> <span class="k">for</span> <span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="ow">in</span> <span class="n">ranges</span><span class="p">[</span><span class="mi">1</span><span class="p">:]:</span>
</span><span id="L-452"><a href="#L-452"><span class="linenos">452</span></a> <span class="n">last_start</span><span class="p">,</span> <span class="n">last_end</span> <span class="o">=</span> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
</span><span id="L-453"><a href="#L-453"><span class="linenos">453</span></a>
</span><span id="L-454"><a href="#L-454"><span class="linenos">454</span></a> <span class="k">if</span> <span class="n">start</span> <span class="o">&lt;=</span> <span class="n">last_end</span><span class="p">:</span>
</span><span id="L-455"><a href="#L-455"><span class="linenos">455</span></a> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">last_start</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="n">last_end</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="L-456"><a href="#L-456"><span class="linenos">456</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-457"><a href="#L-457"><span class="linenos">457</span></a> <span class="n">merged</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="L-444"><a href="#L-444"><span class="linenos">444</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-445"><a href="#L-445"><span class="linenos">445</span></a><span class="sd"> Merges a sequence of ranges, represented as tuples (low, high) whose values</span>
</span><span id="L-446"><a href="#L-446"><span class="linenos">446</span></a><span class="sd"> belong to some totally-ordered set.</span>
</span><span id="L-447"><a href="#L-447"><span class="linenos">447</span></a>
</span><span id="L-448"><a href="#L-448"><span class="linenos">448</span></a><span class="sd"> Example:</span>
</span><span id="L-449"><a href="#L-449"><span class="linenos">449</span></a><span class="sd"> &gt;&gt;&gt; merge_ranges([(1, 3), (2, 6)])</span>
</span><span id="L-450"><a href="#L-450"><span class="linenos">450</span></a><span class="sd"> [(1, 6)]</span>
</span><span id="L-451"><a href="#L-451"><span class="linenos">451</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-452"><a href="#L-452"><span class="linenos">452</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">ranges</span><span class="p">:</span>
</span><span id="L-453"><a href="#L-453"><span class="linenos">453</span></a> <span class="k">return</span> <span class="p">[]</span>
</span><span id="L-454"><a href="#L-454"><span class="linenos">454</span></a>
</span><span id="L-455"><a href="#L-455"><span class="linenos">455</span></a> <span class="n">ranges</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">ranges</span><span class="p">)</span>
</span><span id="L-456"><a href="#L-456"><span class="linenos">456</span></a>
</span><span id="L-457"><a href="#L-457"><span class="linenos">457</span></a> <span class="n">merged</span> <span class="o">=</span> <span class="p">[</span><span class="n">ranges</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span>
</span><span id="L-458"><a href="#L-458"><span class="linenos">458</span></a>
</span><span id="L-459"><a href="#L-459"><span class="linenos">459</span></a> <span class="k">return</span> <span class="n">merged</span>
</span><span id="L-459"><a href="#L-459"><span class="linenos">459</span></a> <span class="k">for</span> <span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="ow">in</span> <span class="n">ranges</span><span class="p">[</span><span class="mi">1</span><span class="p">:]:</span>
</span><span id="L-460"><a href="#L-460"><span class="linenos">460</span></a> <span class="n">last_start</span><span class="p">,</span> <span class="n">last_end</span> <span class="o">=</span> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
</span><span id="L-461"><a href="#L-461"><span class="linenos">461</span></a>
</span><span id="L-462"><a href="#L-462"><span class="linenos">462</span></a> <span class="k">if</span> <span class="n">start</span> <span class="o">&lt;=</span> <span class="n">last_end</span><span class="p">:</span>
</span><span id="L-463"><a href="#L-463"><span class="linenos">463</span></a> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">last_start</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="n">last_end</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="L-464"><a href="#L-464"><span class="linenos">464</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-465"><a href="#L-465"><span class="linenos">465</span></a> <span class="n">merged</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="L-466"><a href="#L-466"><span class="linenos">466</span></a>
</span><span id="L-467"><a href="#L-467"><span class="linenos">467</span></a> <span class="k">return</span> <span class="n">merged</span>
</span></pre></div>
@ -1641,25 +1649,46 @@ type <code>str</code> and <code>bytes</code> are not regarded as iterables.</p>
</div>
<a class="headerlink" href="#merge_ranges"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="merge_ranges-444"><a href="#merge_ranges-444"><span class="linenos">444</span></a><span class="k">def</span> <span class="nf">merge_ranges</span><span class="p">(</span><span class="n">ranges</span><span class="p">:</span> <span class="n">t</span><span class="o">.</span><span class="n">List</span><span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">A</span><span class="p">]])</span> <span class="o">-&gt;</span> <span class="n">t</span><span class="o">.</span><span class="n">List</span><span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">A</span><span class="p">]]:</span>
</span><span id="merge_ranges-445"><a href="#merge_ranges-445"><span class="linenos">445</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">ranges</span><span class="p">:</span>
</span><span id="merge_ranges-446"><a href="#merge_ranges-446"><span class="linenos">446</span></a> <span class="k">return</span> <span class="p">[]</span>
</span><span id="merge_ranges-447"><a href="#merge_ranges-447"><span class="linenos">447</span></a>
</span><span id="merge_ranges-448"><a href="#merge_ranges-448"><span class="linenos">448</span></a> <span class="n">ranges</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">ranges</span><span class="p">)</span>
</span><span id="merge_ranges-449"><a href="#merge_ranges-449"><span class="linenos">449</span></a>
</span><span id="merge_ranges-450"><a href="#merge_ranges-450"><span class="linenos">450</span></a> <span class="n">merged</span> <span class="o">=</span> <span class="p">[</span><span class="n">ranges</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span>
</span><span id="merge_ranges-451"><a href="#merge_ranges-451"><span class="linenos">451</span></a>
</span><span id="merge_ranges-452"><a href="#merge_ranges-452"><span class="linenos">452</span></a> <span class="k">for</span> <span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="ow">in</span> <span class="n">ranges</span><span class="p">[</span><span class="mi">1</span><span class="p">:]:</span>
</span><span id="merge_ranges-453"><a href="#merge_ranges-453"><span class="linenos">453</span></a> <span class="n">last_start</span><span class="p">,</span> <span class="n">last_end</span> <span class="o">=</span> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
</span><span id="merge_ranges-454"><a href="#merge_ranges-454"><span class="linenos">454</span></a>
</span><span id="merge_ranges-455"><a href="#merge_ranges-455"><span class="linenos">455</span></a> <span class="k">if</span> <span class="n">start</span> <span class="o">&lt;=</span> <span class="n">last_end</span><span class="p">:</span>
</span><span id="merge_ranges-456"><a href="#merge_ranges-456"><span class="linenos">456</span></a> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">last_start</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="n">last_end</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="merge_ranges-457"><a href="#merge_ranges-457"><span class="linenos">457</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="merge_ranges-458"><a href="#merge_ranges-458"><span class="linenos">458</span></a> <span class="n">merged</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="merge_ranges-445"><a href="#merge_ranges-445"><span class="linenos">445</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="merge_ranges-446"><a href="#merge_ranges-446"><span class="linenos">446</span></a><span class="sd"> Merges a sequence of ranges, represented as tuples (low, high) whose values</span>
</span><span id="merge_ranges-447"><a href="#merge_ranges-447"><span class="linenos">447</span></a><span class="sd"> belong to some totally-ordered set.</span>
</span><span id="merge_ranges-448"><a href="#merge_ranges-448"><span class="linenos">448</span></a>
</span><span id="merge_ranges-449"><a href="#merge_ranges-449"><span class="linenos">449</span></a><span class="sd"> Example:</span>
</span><span id="merge_ranges-450"><a href="#merge_ranges-450"><span class="linenos">450</span></a><span class="sd"> &gt;&gt;&gt; merge_ranges([(1, 3), (2, 6)])</span>
</span><span id="merge_ranges-451"><a href="#merge_ranges-451"><span class="linenos">451</span></a><span class="sd"> [(1, 6)]</span>
</span><span id="merge_ranges-452"><a href="#merge_ranges-452"><span class="linenos">452</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="merge_ranges-453"><a href="#merge_ranges-453"><span class="linenos">453</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">ranges</span><span class="p">:</span>
</span><span id="merge_ranges-454"><a href="#merge_ranges-454"><span class="linenos">454</span></a> <span class="k">return</span> <span class="p">[]</span>
</span><span id="merge_ranges-455"><a href="#merge_ranges-455"><span class="linenos">455</span></a>
</span><span id="merge_ranges-456"><a href="#merge_ranges-456"><span class="linenos">456</span></a> <span class="n">ranges</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">ranges</span><span class="p">)</span>
</span><span id="merge_ranges-457"><a href="#merge_ranges-457"><span class="linenos">457</span></a>
</span><span id="merge_ranges-458"><a href="#merge_ranges-458"><span class="linenos">458</span></a> <span class="n">merged</span> <span class="o">=</span> <span class="p">[</span><span class="n">ranges</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span>
</span><span id="merge_ranges-459"><a href="#merge_ranges-459"><span class="linenos">459</span></a>
</span><span id="merge_ranges-460"><a href="#merge_ranges-460"><span class="linenos">460</span></a> <span class="k">return</span> <span class="n">merged</span>
</span><span id="merge_ranges-460"><a href="#merge_ranges-460"><span class="linenos">460</span></a> <span class="k">for</span> <span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="ow">in</span> <span class="n">ranges</span><span class="p">[</span><span class="mi">1</span><span class="p">:]:</span>
</span><span id="merge_ranges-461"><a href="#merge_ranges-461"><span class="linenos">461</span></a> <span class="n">last_start</span><span class="p">,</span> <span class="n">last_end</span> <span class="o">=</span> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
</span><span id="merge_ranges-462"><a href="#merge_ranges-462"><span class="linenos">462</span></a>
</span><span id="merge_ranges-463"><a href="#merge_ranges-463"><span class="linenos">463</span></a> <span class="k">if</span> <span class="n">start</span> <span class="o">&lt;=</span> <span class="n">last_end</span><span class="p">:</span>
</span><span id="merge_ranges-464"><a href="#merge_ranges-464"><span class="linenos">464</span></a> <span class="n">merged</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">last_start</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="n">last_end</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="merge_ranges-465"><a href="#merge_ranges-465"><span class="linenos">465</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="merge_ranges-466"><a href="#merge_ranges-466"><span class="linenos">466</span></a> <span class="n">merged</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">))</span>
</span><span id="merge_ranges-467"><a href="#merge_ranges-467"><span class="linenos">467</span></a>
</span><span id="merge_ranges-468"><a href="#merge_ranges-468"><span class="linenos">468</span></a> <span class="k">return</span> <span class="n">merged</span>
</span></pre></div>
<div class="docstring"><p>Merges a sequence of ranges, represented as tuples (low, high) whose values
belong to some totally-ordered set.</p>
<h6 id="example">Example:</h6>
<blockquote>
<div class="pdoc-code codehilite">
<pre><span></span><code><span class="gp">&gt;&gt;&gt; </span><span class="n">merge_ranges</span><span class="p">([(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">6</span><span class="p">)])</span>
<span class="go">[(1, 6)]</span>
</code></pre>
</div>
</blockquote>
</div>
</section>

File diff suppressed because one or more lines are too long

View file

@ -568,7 +568,7 @@ queries if it would result in multiple table selects in a single query:</p>
<div class="attr variable">
<span class="name">UNMERGABLE_ARGS</span> =
<input id="UNMERGABLE_ARGS-view-value" class="view-value-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
<label class="view-value-button pdoc-button" for="UNMERGABLE_ARGS-view-value"></label><span class="default_value">{&#39;into&#39;, &#39;group&#39;, &#39;kind&#39;, &#39;settings&#39;, &#39;having&#39;, &#39;sample&#39;, &#39;match&#39;, &#39;windows&#39;, &#39;distinct&#39;, &#39;connect&#39;, &#39;with&#39;, &#39;format&#39;, &#39;limit&#39;, &#39;offset&#39;, &#39;distribute&#39;, &#39;qualify&#39;, &#39;cluster&#39;, &#39;laterals&#39;, &#39;pivots&#39;, &#39;locks&#39;, &#39;sort&#39;}</span>
<label class="view-value-button pdoc-button" for="UNMERGABLE_ARGS-view-value"></label><span class="default_value">{&#39;locks&#39;, &#39;distribute&#39;, &#39;settings&#39;, &#39;having&#39;, &#39;format&#39;, &#39;pivots&#39;, &#39;match&#39;, &#39;qualify&#39;, &#39;distinct&#39;, &#39;windows&#39;, &#39;with&#39;, &#39;connect&#39;, &#39;group&#39;, &#39;into&#39;, &#39;kind&#39;, &#39;cluster&#39;, &#39;sort&#39;, &#39;laterals&#39;, &#39;limit&#39;, &#39;sample&#39;, &#39;offset&#39;}</span>
</div>

View file

@ -76,164 +76,185 @@
</span><span id="L-6"><a href="#L-6"><span class="linenos"> 6</span></a><span class="kn">from</span> <span class="nn">sqlglot.errors</span> <span class="kn">import</span> <span class="n">OptimizeError</span>
</span><span id="L-7"><a href="#L-7"><span class="linenos"> 7</span></a><span class="kn">from</span> <span class="nn">sqlglot.generator</span> <span class="kn">import</span> <span class="n">cached_generator</span>
</span><span id="L-8"><a href="#L-8"><span class="linenos"> 8</span></a><span class="kn">from</span> <span class="nn">sqlglot.helper</span> <span class="kn">import</span> <span class="n">while_changing</span>
</span><span id="L-9"><a href="#L-9"><span class="linenos"> 9</span></a><span class="kn">from</span> <span class="nn">sqlglot.optimizer.simplify</span> <span class="kn">import</span> <span class="n">flatten</span><span class="p">,</span> <span class="n">rewrite_between</span><span class="p">,</span> <span class="n">uniq_sort</span>
</span><span id="L-10"><a href="#L-10"><span class="linenos"> 10</span></a>
</span><span id="L-11"><a href="#L-11"><span class="linenos"> 11</span></a><span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="s2">&quot;sqlglot&quot;</span><span class="p">)</span>
</span><span id="L-12"><a href="#L-12"><span class="linenos"> 12</span></a>
</span><span id="L-9"><a href="#L-9"><span class="linenos"> 9</span></a><span class="kn">from</span> <span class="nn">sqlglot.optimizer.scope</span> <span class="kn">import</span> <span class="n">find_all_in_scope</span>
</span><span id="L-10"><a href="#L-10"><span class="linenos"> 10</span></a><span class="kn">from</span> <span class="nn">sqlglot.optimizer.simplify</span> <span class="kn">import</span> <span class="n">flatten</span><span class="p">,</span> <span class="n">rewrite_between</span><span class="p">,</span> <span class="n">uniq_sort</span>
</span><span id="L-11"><a href="#L-11"><span class="linenos"> 11</span></a>
</span><span id="L-12"><a href="#L-12"><span class="linenos"> 12</span></a><span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="s2">&quot;sqlglot&quot;</span><span class="p">)</span>
</span><span id="L-13"><a href="#L-13"><span class="linenos"> 13</span></a>
</span><span id="L-14"><a href="#L-14"><span class="linenos"> 14</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">128</span><span class="p">):</span>
</span><span id="L-15"><a href="#L-15"><span class="linenos"> 15</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-16"><a href="#L-16"><span class="linenos"> 16</span></a><span class="sd"> Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form.</span>
</span><span id="L-17"><a href="#L-17"><span class="linenos"> 17</span></a>
</span><span id="L-18"><a href="#L-18"><span class="linenos"> 18</span></a><span class="sd"> Example:</span>
</span><span id="L-19"><a href="#L-19"><span class="linenos"> 19</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="L-20"><a href="#L-20"><span class="linenos"> 20</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(x AND y) OR z&quot;)</span>
</span><span id="L-21"><a href="#L-21"><span class="linenos"> 21</span></a><span class="sd"> &gt;&gt;&gt; normalize(expression, dnf=False).sql()</span>
</span><span id="L-22"><a href="#L-22"><span class="linenos"> 22</span></a><span class="sd"> &#39;(x OR z) AND (y OR z)&#39;</span>
</span><span id="L-23"><a href="#L-23"><span class="linenos"> 23</span></a>
</span><span id="L-24"><a href="#L-24"><span class="linenos"> 24</span></a><span class="sd"> Args:</span>
</span><span id="L-25"><a href="#L-25"><span class="linenos"> 25</span></a><span class="sd"> expression: expression to normalize</span>
</span><span id="L-26"><a href="#L-26"><span class="linenos"> 26</span></a><span class="sd"> dnf: rewrite in disjunctive normal form instead.</span>
</span><span id="L-27"><a href="#L-27"><span class="linenos"> 27</span></a><span class="sd"> max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion</span>
</span><span id="L-28"><a href="#L-28"><span class="linenos"> 28</span></a><span class="sd"> Returns:</span>
</span><span id="L-29"><a href="#L-29"><span class="linenos"> 29</span></a><span class="sd"> sqlglot.Expression: normalized expression</span>
</span><span id="L-30"><a href="#L-30"><span class="linenos"> 30</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-31"><a href="#L-31"><span class="linenos"> 31</span></a> <span class="n">generate</span> <span class="o">=</span> <span class="n">cached_generator</span><span class="p">()</span>
</span><span id="L-32"><a href="#L-32"><span class="linenos"> 32</span></a>
</span><span id="L-33"><a href="#L-33"><span class="linenos"> 33</span></a> <span class="k">for</span> <span class="n">node</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span> <span class="ow">in</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">expression</span><span class="o">.</span><span class="n">walk</span><span class="p">(</span><span class="n">prune</span><span class="o">=</span><span class="k">lambda</span> <span class="n">e</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span><span class="p">:</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="L-34"><a href="#L-34"><span class="linenos"> 34</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="L-35"><a href="#L-35"><span class="linenos"> 35</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="L-36"><a href="#L-36"><span class="linenos"> 36</span></a> <span class="k">continue</span>
</span><span id="L-37"><a href="#L-37"><span class="linenos"> 37</span></a> <span class="n">root</span> <span class="o">=</span> <span class="n">node</span> <span class="ow">is</span> <span class="n">expression</span>
</span><span id="L-38"><a href="#L-38"><span class="linenos"> 38</span></a> <span class="n">original</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
</span><span id="L-39"><a href="#L-39"><span class="linenos"> 39</span></a>
</span><span id="L-40"><a href="#L-40"><span class="linenos"> 40</span></a> <span class="n">node</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">rewrite_between</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-41"><a href="#L-41"><span class="linenos"> 41</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="L-42"><a href="#L-42"><span class="linenos"> 42</span></a>
</span><span id="L-43"><a href="#L-43"><span class="linenos"> 43</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="L-44"><a href="#L-44"><span class="linenos"> 44</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span>
</span><span id="L-45"><a href="#L-45"><span class="linenos"> 45</span></a> <span class="sa">f</span><span class="s2">&quot;Skipping normalization because distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span>
</span><span id="L-46"><a href="#L-46"><span class="linenos"> 46</span></a> <span class="p">)</span>
</span><span id="L-47"><a href="#L-47"><span class="linenos"> 47</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-48"><a href="#L-48"><span class="linenos"> 48</span></a>
</span><span id="L-49"><a href="#L-49"><span class="linenos"> 49</span></a> <span class="k">try</span><span class="p">:</span>
</span><span id="L-50"><a href="#L-50"><span class="linenos"> 50</span></a> <span class="n">node</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span>
</span><span id="L-51"><a href="#L-51"><span class="linenos"> 51</span></a> <span class="n">while_changing</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos"> 52</span></a> <span class="p">)</span>
</span><span id="L-53"><a href="#L-53"><span class="linenos"> 53</span></a> <span class="k">except</span> <span class="n">OptimizeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
</span><span id="L-54"><a href="#L-54"><span class="linenos"> 54</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
</span><span id="L-55"><a href="#L-55"><span class="linenos"> 55</span></a> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">original</span><span class="p">)</span>
</span><span id="L-56"><a href="#L-56"><span class="linenos"> 56</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="L-57"><a href="#L-57"><span class="linenos"> 57</span></a> <span class="k">return</span> <span class="n">original</span>
</span><span id="L-58"><a href="#L-58"><span class="linenos"> 58</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-59"><a href="#L-59"><span class="linenos"> 59</span></a>
</span><span id="L-60"><a href="#L-60"><span class="linenos"> 60</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="L-61"><a href="#L-61"><span class="linenos"> 61</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">node</span>
</span><span id="L-62"><a href="#L-62"><span class="linenos"> 62</span></a>
</span><span id="L-63"><a href="#L-63"><span class="linenos"> 63</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-64"><a href="#L-64"><span class="linenos"> 64</span></a>
</span><span id="L-14"><a href="#L-14"><span class="linenos"> 14</span></a>
</span><span id="L-15"><a href="#L-15"><span class="linenos"> 15</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">128</span><span class="p">):</span>
</span><span id="L-16"><a href="#L-16"><span class="linenos"> 16</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-17"><a href="#L-17"><span class="linenos"> 17</span></a><span class="sd"> Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form.</span>
</span><span id="L-18"><a href="#L-18"><span class="linenos"> 18</span></a>
</span><span id="L-19"><a href="#L-19"><span class="linenos"> 19</span></a><span class="sd"> Example:</span>
</span><span id="L-20"><a href="#L-20"><span class="linenos"> 20</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="L-21"><a href="#L-21"><span class="linenos"> 21</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(x AND y) OR z&quot;)</span>
</span><span id="L-22"><a href="#L-22"><span class="linenos"> 22</span></a><span class="sd"> &gt;&gt;&gt; normalize(expression, dnf=False).sql()</span>
</span><span id="L-23"><a href="#L-23"><span class="linenos"> 23</span></a><span class="sd"> &#39;(x OR z) AND (y OR z)&#39;</span>
</span><span id="L-24"><a href="#L-24"><span class="linenos"> 24</span></a>
</span><span id="L-25"><a href="#L-25"><span class="linenos"> 25</span></a><span class="sd"> Args:</span>
</span><span id="L-26"><a href="#L-26"><span class="linenos"> 26</span></a><span class="sd"> expression: expression to normalize</span>
</span><span id="L-27"><a href="#L-27"><span class="linenos"> 27</span></a><span class="sd"> dnf: rewrite in disjunctive normal form instead.</span>
</span><span id="L-28"><a href="#L-28"><span class="linenos"> 28</span></a><span class="sd"> max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion</span>
</span><span id="L-29"><a href="#L-29"><span class="linenos"> 29</span></a><span class="sd"> Returns:</span>
</span><span id="L-30"><a href="#L-30"><span class="linenos"> 30</span></a><span class="sd"> sqlglot.Expression: normalized expression</span>
</span><span id="L-31"><a href="#L-31"><span class="linenos"> 31</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-32"><a href="#L-32"><span class="linenos"> 32</span></a> <span class="n">generate</span> <span class="o">=</span> <span class="n">cached_generator</span><span class="p">()</span>
</span><span id="L-33"><a href="#L-33"><span class="linenos"> 33</span></a>
</span><span id="L-34"><a href="#L-34"><span class="linenos"> 34</span></a> <span class="k">for</span> <span class="n">node</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span> <span class="ow">in</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">expression</span><span class="o">.</span><span class="n">walk</span><span class="p">(</span><span class="n">prune</span><span class="o">=</span><span class="k">lambda</span> <span class="n">e</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span><span class="p">:</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="L-35"><a href="#L-35"><span class="linenos"> 35</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="L-36"><a href="#L-36"><span class="linenos"> 36</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="L-37"><a href="#L-37"><span class="linenos"> 37</span></a> <span class="k">continue</span>
</span><span id="L-38"><a href="#L-38"><span class="linenos"> 38</span></a> <span class="n">root</span> <span class="o">=</span> <span class="n">node</span> <span class="ow">is</span> <span class="n">expression</span>
</span><span id="L-39"><a href="#L-39"><span class="linenos"> 39</span></a> <span class="n">original</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
</span><span id="L-40"><a href="#L-40"><span class="linenos"> 40</span></a>
</span><span id="L-41"><a href="#L-41"><span class="linenos"> 41</span></a> <span class="n">node</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">rewrite_between</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-42"><a href="#L-42"><span class="linenos"> 42</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="L-43"><a href="#L-43"><span class="linenos"> 43</span></a>
</span><span id="L-44"><a href="#L-44"><span class="linenos"> 44</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="L-45"><a href="#L-45"><span class="linenos"> 45</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span>
</span><span id="L-46"><a href="#L-46"><span class="linenos"> 46</span></a> <span class="sa">f</span><span class="s2">&quot;Skipping normalization because distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span>
</span><span id="L-47"><a href="#L-47"><span class="linenos"> 47</span></a> <span class="p">)</span>
</span><span id="L-48"><a href="#L-48"><span class="linenos"> 48</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-49"><a href="#L-49"><span class="linenos"> 49</span></a>
</span><span id="L-50"><a href="#L-50"><span class="linenos"> 50</span></a> <span class="k">try</span><span class="p">:</span>
</span><span id="L-51"><a href="#L-51"><span class="linenos"> 51</span></a> <span class="n">node</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos"> 52</span></a> <span class="n">while_changing</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="L-53"><a href="#L-53"><span class="linenos"> 53</span></a> <span class="p">)</span>
</span><span id="L-54"><a href="#L-54"><span class="linenos"> 54</span></a> <span class="k">except</span> <span class="n">OptimizeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
</span><span id="L-55"><a href="#L-55"><span class="linenos"> 55</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
</span><span id="L-56"><a href="#L-56"><span class="linenos"> 56</span></a> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">original</span><span class="p">)</span>
</span><span id="L-57"><a href="#L-57"><span class="linenos"> 57</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="L-58"><a href="#L-58"><span class="linenos"> 58</span></a> <span class="k">return</span> <span class="n">original</span>
</span><span id="L-59"><a href="#L-59"><span class="linenos"> 59</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-60"><a href="#L-60"><span class="linenos"> 60</span></a>
</span><span id="L-61"><a href="#L-61"><span class="linenos"> 61</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="L-62"><a href="#L-62"><span class="linenos"> 62</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">node</span>
</span><span id="L-63"><a href="#L-63"><span class="linenos"> 63</span></a>
</span><span id="L-64"><a href="#L-64"><span class="linenos"> 64</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-65"><a href="#L-65"><span class="linenos"> 65</span></a>
</span><span id="L-66"><a href="#L-66"><span class="linenos"> 66</span></a><span class="k">def</span> <span class="nf">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
</span><span id="L-67"><a href="#L-67"><span class="linenos"> 67</span></a> <span class="n">ancestor</span><span class="p">,</span> <span class="n">root</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span>
</span><span id="L-68"><a href="#L-68"><span class="linenos"> 68</span></a>
</span><span id="L-69"><a href="#L-69"><span class="linenos"> 69</span></a> <span class="k">return</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">connector</span><span class="o">.</span><span class="n">find_ancestor</span><span class="p">(</span><span class="n">ancestor</span><span class="p">)</span> <span class="k">for</span> <span class="n">connector</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">root</span><span class="p">))</span>
</span><span id="L-66"><a href="#L-66"><span class="linenos"> 66</span></a>
</span><span id="L-67"><a href="#L-67"><span class="linenos"> 67</span></a><span class="k">def</span> <span class="nf">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
</span><span id="L-68"><a href="#L-68"><span class="linenos"> 68</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-69"><a href="#L-69"><span class="linenos"> 69</span></a><span class="sd"> Checks whether a given expression is in a normal form of interest.</span>
</span><span id="L-70"><a href="#L-70"><span class="linenos"> 70</span></a>
</span><span id="L-71"><a href="#L-71"><span class="linenos"> 71</span></a>
</span><span id="L-72"><a href="#L-72"><span class="linenos"> 72</span></a><span class="k">def</span> <span class="nf">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
</span><span id="L-73"><a href="#L-73"><span class="linenos"> 73</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-74"><a href="#L-74"><span class="linenos"> 74</span></a><span class="sd"> The difference in the number of predicates between the current expression and the normalized form.</span>
</span><span id="L-75"><a href="#L-75"><span class="linenos"> 75</span></a>
</span><span id="L-76"><a href="#L-76"><span class="linenos"> 76</span></a><span class="sd"> This is used as an estimate of the cost of the conversion which is exponential in complexity.</span>
</span><span id="L-77"><a href="#L-77"><span class="linenos"> 77</span></a>
</span><span id="L-78"><a href="#L-78"><span class="linenos"> 78</span></a><span class="sd"> Example:</span>
</span><span id="L-79"><a href="#L-79"><span class="linenos"> 79</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="L-80"><a href="#L-80"><span class="linenos"> 80</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(a AND b) OR (c AND d)&quot;)</span>
</span><span id="L-81"><a href="#L-81"><span class="linenos"> 81</span></a><span class="sd"> &gt;&gt;&gt; normalization_distance(expression)</span>
</span><span id="L-82"><a href="#L-82"><span class="linenos"> 82</span></a><span class="sd"> 4</span>
</span><span id="L-83"><a href="#L-83"><span class="linenos"> 83</span></a>
</span><span id="L-84"><a href="#L-84"><span class="linenos"> 84</span></a><span class="sd"> Args:</span>
</span><span id="L-85"><a href="#L-85"><span class="linenos"> 85</span></a><span class="sd"> expression (sqlglot.Expression): expression to compute distance</span>
</span><span id="L-86"><a href="#L-86"><span class="linenos"> 86</span></a><span class="sd"> dnf (bool): compute to dnf distance instead</span>
</span><span id="L-87"><a href="#L-87"><span class="linenos"> 87</span></a><span class="sd"> Returns:</span>
</span><span id="L-88"><a href="#L-88"><span class="linenos"> 88</span></a><span class="sd"> int: difference</span>
</span><span id="L-89"><a href="#L-89"><span class="linenos"> 89</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-90"><a href="#L-90"><span class="linenos"> 90</span></a> <span class="k">return</span> <span class="nb">sum</span><span class="p">(</span><span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">))</span> <span class="o">-</span> <span class="p">(</span>
</span><span id="L-91"><a href="#L-91"><span class="linenos"> 91</span></a> <span class="nb">sum</span><span class="p">(</span><span class="mi">1</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))</span> <span class="o">+</span> <span class="mi">1</span>
</span><span id="L-92"><a href="#L-92"><span class="linenos"> 92</span></a> <span class="p">)</span>
</span><span id="L-93"><a href="#L-93"><span class="linenos"> 93</span></a>
</span><span id="L-71"><a href="#L-71"><span class="linenos"> 71</span></a><span class="sd"> Example:</span>
</span><span id="L-72"><a href="#L-72"><span class="linenos"> 72</span></a><span class="sd"> &gt;&gt;&gt; from sqlglot import parse_one</span>
</span><span id="L-73"><a href="#L-73"><span class="linenos"> 73</span></a><span class="sd"> &gt;&gt;&gt; normalized(parse_one(&quot;(a AND b) OR c OR (d AND e)&quot;), dnf=True)</span>
</span><span id="L-74"><a href="#L-74"><span class="linenos"> 74</span></a><span class="sd"> True</span>
</span><span id="L-75"><a href="#L-75"><span class="linenos"> 75</span></a><span class="sd"> &gt;&gt;&gt; normalized(parse_one(&quot;(a OR b) AND c&quot;)) # Checks CNF by default</span>
</span><span id="L-76"><a href="#L-76"><span class="linenos"> 76</span></a><span class="sd"> True</span>
</span><span id="L-77"><a href="#L-77"><span class="linenos"> 77</span></a><span class="sd"> &gt;&gt;&gt; normalized(parse_one(&quot;a AND (b OR c)&quot;), dnf=True)</span>
</span><span id="L-78"><a href="#L-78"><span class="linenos"> 78</span></a><span class="sd"> False</span>
</span><span id="L-79"><a href="#L-79"><span class="linenos"> 79</span></a>
</span><span id="L-80"><a href="#L-80"><span class="linenos"> 80</span></a><span class="sd"> Args:</span>
</span><span id="L-81"><a href="#L-81"><span class="linenos"> 81</span></a><span class="sd"> expression: The expression to check if it&#39;s normalized.</span>
</span><span id="L-82"><a href="#L-82"><span class="linenos"> 82</span></a><span class="sd"> dnf: Whether or not to check if the expression is in Disjunctive Normal Form (DNF).</span>
</span><span id="L-83"><a href="#L-83"><span class="linenos"> 83</span></a><span class="sd"> Default: False, i.e. we check if it&#39;s in Conjunctive Normal Form (CNF).</span>
</span><span id="L-84"><a href="#L-84"><span class="linenos"> 84</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-85"><a href="#L-85"><span class="linenos"> 85</span></a> <span class="n">ancestor</span><span class="p">,</span> <span class="n">root</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span>
</span><span id="L-86"><a href="#L-86"><span class="linenos"> 86</span></a> <span class="k">return</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span>
</span><span id="L-87"><a href="#L-87"><span class="linenos"> 87</span></a> <span class="n">connector</span><span class="o">.</span><span class="n">find_ancestor</span><span class="p">(</span><span class="n">ancestor</span><span class="p">)</span> <span class="k">for</span> <span class="n">connector</span> <span class="ow">in</span> <span class="n">find_all_in_scope</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">root</span><span class="p">)</span>
</span><span id="L-88"><a href="#L-88"><span class="linenos"> 88</span></a> <span class="p">)</span>
</span><span id="L-89"><a href="#L-89"><span class="linenos"> 89</span></a>
</span><span id="L-90"><a href="#L-90"><span class="linenos"> 90</span></a>
</span><span id="L-91"><a href="#L-91"><span class="linenos"> 91</span></a><span class="k">def</span> <span class="nf">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
</span><span id="L-92"><a href="#L-92"><span class="linenos"> 92</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-93"><a href="#L-93"><span class="linenos"> 93</span></a><span class="sd"> The difference in the number of predicates between a given expression and its normalized form.</span>
</span><span id="L-94"><a href="#L-94"><span class="linenos"> 94</span></a>
</span><span id="L-95"><a href="#L-95"><span class="linenos"> 95</span></a><span class="k">def</span> <span class="nf">_predicate_lengths</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">):</span>
</span><span id="L-96"><a href="#L-96"><span class="linenos"> 96</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-97"><a href="#L-97"><span class="linenos"> 97</span></a><span class="sd"> Returns a list of predicate lengths when expanded to normalized form.</span>
</span><span id="L-98"><a href="#L-98"><span class="linenos"> 98</span></a>
</span><span id="L-99"><a href="#L-99"><span class="linenos"> 99</span></a><span class="sd"> (A AND B) OR C -&gt; [2, 2] because len(A OR C), len(B OR C).</span>
</span><span id="L-100"><a href="#L-100"><span class="linenos">100</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-101"><a href="#L-101"><span class="linenos">101</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">unnest</span><span class="p">()</span>
</span><span id="L-95"><a href="#L-95"><span class="linenos"> 95</span></a><span class="sd"> This is used as an estimate of the cost of the conversion which is exponential in complexity.</span>
</span><span id="L-96"><a href="#L-96"><span class="linenos"> 96</span></a>
</span><span id="L-97"><a href="#L-97"><span class="linenos"> 97</span></a><span class="sd"> Example:</span>
</span><span id="L-98"><a href="#L-98"><span class="linenos"> 98</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="L-99"><a href="#L-99"><span class="linenos"> 99</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(a AND b) OR (c AND d)&quot;)</span>
</span><span id="L-100"><a href="#L-100"><span class="linenos">100</span></a><span class="sd"> &gt;&gt;&gt; normalization_distance(expression)</span>
</span><span id="L-101"><a href="#L-101"><span class="linenos">101</span></a><span class="sd"> 4</span>
</span><span id="L-102"><a href="#L-102"><span class="linenos">102</span></a>
</span><span id="L-103"><a href="#L-103"><span class="linenos">103</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="L-104"><a href="#L-104"><span class="linenos">104</span></a> <span class="k">return</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span>
</span><span id="L-105"><a href="#L-105"><span class="linenos">105</span></a>
</span><span id="L-106"><a href="#L-106"><span class="linenos">106</span></a> <span class="n">left</span><span class="p">,</span> <span class="n">right</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">values</span><span class="p">()</span>
</span><span id="L-103"><a href="#L-103"><span class="linenos">103</span></a><span class="sd"> Args:</span>
</span><span id="L-104"><a href="#L-104"><span class="linenos">104</span></a><span class="sd"> expression: The expression to compute the normalization distance for.</span>
</span><span id="L-105"><a href="#L-105"><span class="linenos">105</span></a><span class="sd"> dnf: Whether or not to check if the expression is in Disjunctive Normal Form (DNF).</span>
</span><span id="L-106"><a href="#L-106"><span class="linenos">106</span></a><span class="sd"> Default: False, i.e. we check if it&#39;s in Conjunctive Normal Form (CNF).</span>
</span><span id="L-107"><a href="#L-107"><span class="linenos">107</span></a>
</span><span id="L-108"><a href="#L-108"><span class="linenos">108</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">):</span>
</span><span id="L-109"><a href="#L-109"><span class="linenos">109</span></a> <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span>
</span><span id="L-110"><a href="#L-110"><span class="linenos">110</span></a> <span class="n">a</span> <span class="o">+</span> <span class="n">b</span> <span class="k">for</span> <span class="n">a</span> <span class="ow">in</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span> <span class="k">for</span> <span class="n">b</span> <span class="ow">in</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span>
</span><span id="L-111"><a href="#L-111"><span class="linenos">111</span></a> <span class="p">)</span>
</span><span id="L-112"><a href="#L-112"><span class="linenos">112</span></a> <span class="k">return</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span> <span class="o">+</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span>
</span><span id="L-113"><a href="#L-113"><span class="linenos">113</span></a>
</span><span id="L-108"><a href="#L-108"><span class="linenos">108</span></a><span class="sd"> Returns:</span>
</span><span id="L-109"><a href="#L-109"><span class="linenos">109</span></a><span class="sd"> The normalization distance.</span>
</span><span id="L-110"><a href="#L-110"><span class="linenos">110</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-111"><a href="#L-111"><span class="linenos">111</span></a> <span class="k">return</span> <span class="nb">sum</span><span class="p">(</span><span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">))</span> <span class="o">-</span> <span class="p">(</span>
</span><span id="L-112"><a href="#L-112"><span class="linenos">112</span></a> <span class="nb">sum</span><span class="p">(</span><span class="mi">1</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))</span> <span class="o">+</span> <span class="mi">1</span>
</span><span id="L-113"><a href="#L-113"><span class="linenos">113</span></a> <span class="p">)</span>
</span><span id="L-114"><a href="#L-114"><span class="linenos">114</span></a>
</span><span id="L-115"><a href="#L-115"><span class="linenos">115</span></a><span class="k">def</span> <span class="nf">distributive_law</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">):</span>
</span><span id="L-116"><a href="#L-116"><span class="linenos">116</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-117"><a href="#L-117"><span class="linenos">117</span></a><span class="sd"> x OR (y AND z) -&gt; (x OR y) AND (x OR z)</span>
</span><span id="L-118"><a href="#L-118"><span class="linenos">118</span></a><span class="sd"> (x AND y) OR (y AND z) -&gt; (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)</span>
</span><span id="L-119"><a href="#L-119"><span class="linenos">119</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-120"><a href="#L-120"><span class="linenos">120</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="L-121"><a href="#L-121"><span class="linenos">121</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-122"><a href="#L-122"><span class="linenos">122</span></a>
</span><span id="L-123"><a href="#L-123"><span class="linenos">123</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="L-124"><a href="#L-124"><span class="linenos">124</span></a>
</span><span id="L-125"><a href="#L-125"><span class="linenos">125</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="L-126"><a href="#L-126"><span class="linenos">126</span></a> <span class="k">raise</span> <span class="n">OptimizeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Normalization distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
</span><span id="L-127"><a href="#L-127"><span class="linenos">127</span></a>
</span><span id="L-128"><a href="#L-128"><span class="linenos">128</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">replace_children</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="L-129"><a href="#L-129"><span class="linenos">129</span></a> <span class="n">to_exp</span><span class="p">,</span> <span class="n">from_exp</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span>
</span><span id="L-130"><a href="#L-130"><span class="linenos">130</span></a>
</span><span id="L-131"><a href="#L-131"><span class="linenos">131</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">from_exp</span><span class="p">):</span>
</span><span id="L-132"><a href="#L-132"><span class="linenos">132</span></a> <span class="n">a</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">unnest_operands</span><span class="p">()</span>
</span><span id="L-133"><a href="#L-133"><span class="linenos">133</span></a>
</span><span id="L-134"><a href="#L-134"><span class="linenos">134</span></a> <span class="n">from_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">from_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="L-135"><a href="#L-135"><span class="linenos">135</span></a> <span class="n">to_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">to_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="L-136"><a href="#L-136"><span class="linenos">136</span></a>
</span><span id="L-137"><a href="#L-137"><span class="linenos">137</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="L-138"><a href="#L-138"><span class="linenos">138</span></a> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">)))</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">b</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="L-139"><a href="#L-139"><span class="linenos">139</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-140"><a href="#L-140"><span class="linenos">140</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-141"><a href="#L-141"><span class="linenos">141</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="L-142"><a href="#L-142"><span class="linenos">142</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-143"><a href="#L-143"><span class="linenos">143</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="L-144"><a href="#L-144"><span class="linenos">144</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-115"><a href="#L-115"><span class="linenos">115</span></a>
</span><span id="L-116"><a href="#L-116"><span class="linenos">116</span></a><span class="k">def</span> <span class="nf">_predicate_lengths</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">):</span>
</span><span id="L-117"><a href="#L-117"><span class="linenos">117</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-118"><a href="#L-118"><span class="linenos">118</span></a><span class="sd"> Returns a list of predicate lengths when expanded to normalized form.</span>
</span><span id="L-119"><a href="#L-119"><span class="linenos">119</span></a>
</span><span id="L-120"><a href="#L-120"><span class="linenos">120</span></a><span class="sd"> (A AND B) OR C -&gt; [2, 2] because len(A OR C), len(B OR C).</span>
</span><span id="L-121"><a href="#L-121"><span class="linenos">121</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-122"><a href="#L-122"><span class="linenos">122</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">unnest</span><span class="p">()</span>
</span><span id="L-123"><a href="#L-123"><span class="linenos">123</span></a>
</span><span id="L-124"><a href="#L-124"><span class="linenos">124</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="L-125"><a href="#L-125"><span class="linenos">125</span></a> <span class="k">return</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span>
</span><span id="L-126"><a href="#L-126"><span class="linenos">126</span></a>
</span><span id="L-127"><a href="#L-127"><span class="linenos">127</span></a> <span class="n">left</span><span class="p">,</span> <span class="n">right</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">values</span><span class="p">()</span>
</span><span id="L-128"><a href="#L-128"><span class="linenos">128</span></a>
</span><span id="L-129"><a href="#L-129"><span class="linenos">129</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">):</span>
</span><span id="L-130"><a href="#L-130"><span class="linenos">130</span></a> <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span>
</span><span id="L-131"><a href="#L-131"><span class="linenos">131</span></a> <span class="n">a</span> <span class="o">+</span> <span class="n">b</span> <span class="k">for</span> <span class="n">a</span> <span class="ow">in</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span> <span class="k">for</span> <span class="n">b</span> <span class="ow">in</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span>
</span><span id="L-132"><a href="#L-132"><span class="linenos">132</span></a> <span class="p">)</span>
</span><span id="L-133"><a href="#L-133"><span class="linenos">133</span></a> <span class="k">return</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span> <span class="o">+</span> <span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">dnf</span><span class="p">)</span>
</span><span id="L-134"><a href="#L-134"><span class="linenos">134</span></a>
</span><span id="L-135"><a href="#L-135"><span class="linenos">135</span></a>
</span><span id="L-136"><a href="#L-136"><span class="linenos">136</span></a><span class="k">def</span> <span class="nf">distributive_law</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">):</span>
</span><span id="L-137"><a href="#L-137"><span class="linenos">137</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-138"><a href="#L-138"><span class="linenos">138</span></a><span class="sd"> x OR (y AND z) -&gt; (x OR y) AND (x OR z)</span>
</span><span id="L-139"><a href="#L-139"><span class="linenos">139</span></a><span class="sd"> (x AND y) OR (y AND z) -&gt; (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)</span>
</span><span id="L-140"><a href="#L-140"><span class="linenos">140</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-141"><a href="#L-141"><span class="linenos">141</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="L-142"><a href="#L-142"><span class="linenos">142</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-143"><a href="#L-143"><span class="linenos">143</span></a>
</span><span id="L-144"><a href="#L-144"><span class="linenos">144</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="L-145"><a href="#L-145"><span class="linenos">145</span></a>
</span><span id="L-146"><a href="#L-146"><span class="linenos">146</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-147"><a href="#L-147"><span class="linenos">147</span></a>
</span><span id="L-146"><a href="#L-146"><span class="linenos">146</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="L-147"><a href="#L-147"><span class="linenos">147</span></a> <span class="k">raise</span> <span class="n">OptimizeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Normalization distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
</span><span id="L-148"><a href="#L-148"><span class="linenos">148</span></a>
</span><span id="L-149"><a href="#L-149"><span class="linenos">149</span></a><span class="k">def</span> <span class="nf">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">):</span>
</span><span id="L-150"><a href="#L-150"><span class="linenos">150</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="L-151"><a href="#L-151"><span class="linenos">151</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">replace_children</span><span class="p">(</span>
</span><span id="L-152"><a href="#L-152"><span class="linenos">152</span></a> <span class="n">a</span><span class="p">,</span>
</span><span id="L-153"><a href="#L-153"><span class="linenos">153</span></a> <span class="k">lambda</span> <span class="n">c</span><span class="p">:</span> <span class="n">to_func</span><span class="p">(</span>
</span><span id="L-154"><a href="#L-154"><span class="linenos">154</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">c</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">left</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-155"><a href="#L-155"><span class="linenos">155</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">c</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">right</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-156"><a href="#L-156"><span class="linenos">156</span></a> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span><span id="L-157"><a href="#L-157"><span class="linenos">157</span></a> <span class="p">),</span>
</span><span id="L-158"><a href="#L-158"><span class="linenos">158</span></a> <span class="p">)</span>
</span><span id="L-159"><a href="#L-159"><span class="linenos">159</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-160"><a href="#L-160"><span class="linenos">160</span></a> <span class="n">a</span> <span class="o">=</span> <span class="n">to_func</span><span class="p">(</span>
</span><span id="L-161"><a href="#L-161"><span class="linenos">161</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">left</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-162"><a href="#L-162"><span class="linenos">162</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">right</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-163"><a href="#L-163"><span class="linenos">163</span></a> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span><span id="L-164"><a href="#L-164"><span class="linenos">164</span></a> <span class="p">)</span>
</span><span id="L-165"><a href="#L-165"><span class="linenos">165</span></a>
</span><span id="L-166"><a href="#L-166"><span class="linenos">166</span></a> <span class="k">return</span> <span class="n">a</span>
</span><span id="L-149"><a href="#L-149"><span class="linenos">149</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">replace_children</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="L-150"><a href="#L-150"><span class="linenos">150</span></a> <span class="n">to_exp</span><span class="p">,</span> <span class="n">from_exp</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span>
</span><span id="L-151"><a href="#L-151"><span class="linenos">151</span></a>
</span><span id="L-152"><a href="#L-152"><span class="linenos">152</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">from_exp</span><span class="p">):</span>
</span><span id="L-153"><a href="#L-153"><span class="linenos">153</span></a> <span class="n">a</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">unnest_operands</span><span class="p">()</span>
</span><span id="L-154"><a href="#L-154"><span class="linenos">154</span></a>
</span><span id="L-155"><a href="#L-155"><span class="linenos">155</span></a> <span class="n">from_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">from_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="L-156"><a href="#L-156"><span class="linenos">156</span></a> <span class="n">to_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">to_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="L-157"><a href="#L-157"><span class="linenos">157</span></a>
</span><span id="L-158"><a href="#L-158"><span class="linenos">158</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="L-159"><a href="#L-159"><span class="linenos">159</span></a> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">)))</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">b</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="L-160"><a href="#L-160"><span class="linenos">160</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-161"><a href="#L-161"><span class="linenos">161</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-162"><a href="#L-162"><span class="linenos">162</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="L-163"><a href="#L-163"><span class="linenos">163</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-164"><a href="#L-164"><span class="linenos">164</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="L-165"><a href="#L-165"><span class="linenos">165</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="L-166"><a href="#L-166"><span class="linenos">166</span></a>
</span><span id="L-167"><a href="#L-167"><span class="linenos">167</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-168"><a href="#L-168"><span class="linenos">168</span></a>
</span><span id="L-169"><a href="#L-169"><span class="linenos">169</span></a>
</span><span id="L-170"><a href="#L-170"><span class="linenos">170</span></a><span class="k">def</span> <span class="nf">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">):</span>
</span><span id="L-171"><a href="#L-171"><span class="linenos">171</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="L-172"><a href="#L-172"><span class="linenos">172</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">replace_children</span><span class="p">(</span>
</span><span id="L-173"><a href="#L-173"><span class="linenos">173</span></a> <span class="n">a</span><span class="p">,</span>
</span><span id="L-174"><a href="#L-174"><span class="linenos">174</span></a> <span class="k">lambda</span> <span class="n">c</span><span class="p">:</span> <span class="n">to_func</span><span class="p">(</span>
</span><span id="L-175"><a href="#L-175"><span class="linenos">175</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">c</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">left</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-176"><a href="#L-176"><span class="linenos">176</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">c</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">right</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-177"><a href="#L-177"><span class="linenos">177</span></a> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span><span id="L-178"><a href="#L-178"><span class="linenos">178</span></a> <span class="p">),</span>
</span><span id="L-179"><a href="#L-179"><span class="linenos">179</span></a> <span class="p">)</span>
</span><span id="L-180"><a href="#L-180"><span class="linenos">180</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-181"><a href="#L-181"><span class="linenos">181</span></a> <span class="n">a</span> <span class="o">=</span> <span class="n">to_func</span><span class="p">(</span>
</span><span id="L-182"><a href="#L-182"><span class="linenos">182</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">left</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-183"><a href="#L-183"><span class="linenos">183</span></a> <span class="n">uniq_sort</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">from_func</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">right</span><span class="p">)),</span> <span class="n">generate</span><span class="p">),</span>
</span><span id="L-184"><a href="#L-184"><span class="linenos">184</span></a> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
</span><span id="L-185"><a href="#L-185"><span class="linenos">185</span></a> <span class="p">)</span>
</span><span id="L-186"><a href="#L-186"><span class="linenos">186</span></a>
</span><span id="L-187"><a href="#L-187"><span class="linenos">187</span></a> <span class="k">return</span> <span class="n">a</span>
</span></pre></div>
@ -261,56 +282,56 @@
</div>
<a class="headerlink" href="#normalize"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalize-15"><a href="#normalize-15"><span class="linenos">15</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">128</span><span class="p">):</span>
</span><span id="normalize-16"><a href="#normalize-16"><span class="linenos">16</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalize-17"><a href="#normalize-17"><span class="linenos">17</span></a><span class="sd"> Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form.</span>
</span><span id="normalize-18"><a href="#normalize-18"><span class="linenos">18</span></a>
</span><span id="normalize-19"><a href="#normalize-19"><span class="linenos">19</span></a><span class="sd"> Example:</span>
</span><span id="normalize-20"><a href="#normalize-20"><span class="linenos">20</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="normalize-21"><a href="#normalize-21"><span class="linenos">21</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(x AND y) OR z&quot;)</span>
</span><span id="normalize-22"><a href="#normalize-22"><span class="linenos">22</span></a><span class="sd"> &gt;&gt;&gt; normalize(expression, dnf=False).sql()</span>
</span><span id="normalize-23"><a href="#normalize-23"><span class="linenos">23</span></a><span class="sd"> &#39;(x OR z) AND (y OR z)&#39;</span>
</span><span id="normalize-24"><a href="#normalize-24"><span class="linenos">24</span></a>
</span><span id="normalize-25"><a href="#normalize-25"><span class="linenos">25</span></a><span class="sd"> Args:</span>
</span><span id="normalize-26"><a href="#normalize-26"><span class="linenos">26</span></a><span class="sd"> expression: expression to normalize</span>
</span><span id="normalize-27"><a href="#normalize-27"><span class="linenos">27</span></a><span class="sd"> dnf: rewrite in disjunctive normal form instead.</span>
</span><span id="normalize-28"><a href="#normalize-28"><span class="linenos">28</span></a><span class="sd"> max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion</span>
</span><span id="normalize-29"><a href="#normalize-29"><span class="linenos">29</span></a><span class="sd"> Returns:</span>
</span><span id="normalize-30"><a href="#normalize-30"><span class="linenos">30</span></a><span class="sd"> sqlglot.Expression: normalized expression</span>
</span><span id="normalize-31"><a href="#normalize-31"><span class="linenos">31</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalize-32"><a href="#normalize-32"><span class="linenos">32</span></a> <span class="n">generate</span> <span class="o">=</span> <span class="n">cached_generator</span><span class="p">()</span>
</span><span id="normalize-33"><a href="#normalize-33"><span class="linenos">33</span></a>
</span><span id="normalize-34"><a href="#normalize-34"><span class="linenos">34</span></a> <span class="k">for</span> <span class="n">node</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span> <span class="ow">in</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">expression</span><span class="o">.</span><span class="n">walk</span><span class="p">(</span><span class="n">prune</span><span class="o">=</span><span class="k">lambda</span> <span class="n">e</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span><span class="p">:</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="normalize-35"><a href="#normalize-35"><span class="linenos">35</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="normalize-36"><a href="#normalize-36"><span class="linenos">36</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="normalize-37"><a href="#normalize-37"><span class="linenos">37</span></a> <span class="k">continue</span>
</span><span id="normalize-38"><a href="#normalize-38"><span class="linenos">38</span></a> <span class="n">root</span> <span class="o">=</span> <span class="n">node</span> <span class="ow">is</span> <span class="n">expression</span>
</span><span id="normalize-39"><a href="#normalize-39"><span class="linenos">39</span></a> <span class="n">original</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
</span><span id="normalize-40"><a href="#normalize-40"><span class="linenos">40</span></a>
</span><span id="normalize-41"><a href="#normalize-41"><span class="linenos">41</span></a> <span class="n">node</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">rewrite_between</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="normalize-42"><a href="#normalize-42"><span class="linenos">42</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="normalize-43"><a href="#normalize-43"><span class="linenos">43</span></a>
</span><span id="normalize-44"><a href="#normalize-44"><span class="linenos">44</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="normalize-45"><a href="#normalize-45"><span class="linenos">45</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span>
</span><span id="normalize-46"><a href="#normalize-46"><span class="linenos">46</span></a> <span class="sa">f</span><span class="s2">&quot;Skipping normalization because distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span>
</span><span id="normalize-47"><a href="#normalize-47"><span class="linenos">47</span></a> <span class="p">)</span>
</span><span id="normalize-48"><a href="#normalize-48"><span class="linenos">48</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="normalize-49"><a href="#normalize-49"><span class="linenos">49</span></a>
</span><span id="normalize-50"><a href="#normalize-50"><span class="linenos">50</span></a> <span class="k">try</span><span class="p">:</span>
</span><span id="normalize-51"><a href="#normalize-51"><span class="linenos">51</span></a> <span class="n">node</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span>
</span><span id="normalize-52"><a href="#normalize-52"><span class="linenos">52</span></a> <span class="n">while_changing</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="normalize-53"><a href="#normalize-53"><span class="linenos">53</span></a> <span class="p">)</span>
</span><span id="normalize-54"><a href="#normalize-54"><span class="linenos">54</span></a> <span class="k">except</span> <span class="n">OptimizeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
</span><span id="normalize-55"><a href="#normalize-55"><span class="linenos">55</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
</span><span id="normalize-56"><a href="#normalize-56"><span class="linenos">56</span></a> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">original</span><span class="p">)</span>
</span><span id="normalize-57"><a href="#normalize-57"><span class="linenos">57</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="normalize-58"><a href="#normalize-58"><span class="linenos">58</span></a> <span class="k">return</span> <span class="n">original</span>
</span><span id="normalize-59"><a href="#normalize-59"><span class="linenos">59</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="normalize-60"><a href="#normalize-60"><span class="linenos">60</span></a>
</span><span id="normalize-61"><a href="#normalize-61"><span class="linenos">61</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="normalize-62"><a href="#normalize-62"><span class="linenos">62</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">node</span>
</span><span id="normalize-63"><a href="#normalize-63"><span class="linenos">63</span></a>
</span><span id="normalize-64"><a href="#normalize-64"><span class="linenos">64</span></a> <span class="k">return</span> <span class="n">expression</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalize-16"><a href="#normalize-16"><span class="linenos">16</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">128</span><span class="p">):</span>
</span><span id="normalize-17"><a href="#normalize-17"><span class="linenos">17</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalize-18"><a href="#normalize-18"><span class="linenos">18</span></a><span class="sd"> Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form.</span>
</span><span id="normalize-19"><a href="#normalize-19"><span class="linenos">19</span></a>
</span><span id="normalize-20"><a href="#normalize-20"><span class="linenos">20</span></a><span class="sd"> Example:</span>
</span><span id="normalize-21"><a href="#normalize-21"><span class="linenos">21</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="normalize-22"><a href="#normalize-22"><span class="linenos">22</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(x AND y) OR z&quot;)</span>
</span><span id="normalize-23"><a href="#normalize-23"><span class="linenos">23</span></a><span class="sd"> &gt;&gt;&gt; normalize(expression, dnf=False).sql()</span>
</span><span id="normalize-24"><a href="#normalize-24"><span class="linenos">24</span></a><span class="sd"> &#39;(x OR z) AND (y OR z)&#39;</span>
</span><span id="normalize-25"><a href="#normalize-25"><span class="linenos">25</span></a>
</span><span id="normalize-26"><a href="#normalize-26"><span class="linenos">26</span></a><span class="sd"> Args:</span>
</span><span id="normalize-27"><a href="#normalize-27"><span class="linenos">27</span></a><span class="sd"> expression: expression to normalize</span>
</span><span id="normalize-28"><a href="#normalize-28"><span class="linenos">28</span></a><span class="sd"> dnf: rewrite in disjunctive normal form instead.</span>
</span><span id="normalize-29"><a href="#normalize-29"><span class="linenos">29</span></a><span class="sd"> max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion</span>
</span><span id="normalize-30"><a href="#normalize-30"><span class="linenos">30</span></a><span class="sd"> Returns:</span>
</span><span id="normalize-31"><a href="#normalize-31"><span class="linenos">31</span></a><span class="sd"> sqlglot.Expression: normalized expression</span>
</span><span id="normalize-32"><a href="#normalize-32"><span class="linenos">32</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalize-33"><a href="#normalize-33"><span class="linenos">33</span></a> <span class="n">generate</span> <span class="o">=</span> <span class="n">cached_generator</span><span class="p">()</span>
</span><span id="normalize-34"><a href="#normalize-34"><span class="linenos">34</span></a>
</span><span id="normalize-35"><a href="#normalize-35"><span class="linenos">35</span></a> <span class="k">for</span> <span class="n">node</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span> <span class="ow">in</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">expression</span><span class="o">.</span><span class="n">walk</span><span class="p">(</span><span class="n">prune</span><span class="o">=</span><span class="k">lambda</span> <span class="n">e</span><span class="p">,</span> <span class="o">*</span><span class="n">_</span><span class="p">:</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="normalize-36"><a href="#normalize-36"><span class="linenos">36</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">):</span>
</span><span id="normalize-37"><a href="#normalize-37"><span class="linenos">37</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="normalize-38"><a href="#normalize-38"><span class="linenos">38</span></a> <span class="k">continue</span>
</span><span id="normalize-39"><a href="#normalize-39"><span class="linenos">39</span></a> <span class="n">root</span> <span class="o">=</span> <span class="n">node</span> <span class="ow">is</span> <span class="n">expression</span>
</span><span id="normalize-40"><a href="#normalize-40"><span class="linenos">40</span></a> <span class="n">original</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
</span><span id="normalize-41"><a href="#normalize-41"><span class="linenos">41</span></a>
</span><span id="normalize-42"><a href="#normalize-42"><span class="linenos">42</span></a> <span class="n">node</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">rewrite_between</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="normalize-43"><a href="#normalize-43"><span class="linenos">43</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="normalize-44"><a href="#normalize-44"><span class="linenos">44</span></a>
</span><span id="normalize-45"><a href="#normalize-45"><span class="linenos">45</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="normalize-46"><a href="#normalize-46"><span class="linenos">46</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span>
</span><span id="normalize-47"><a href="#normalize-47"><span class="linenos">47</span></a> <span class="sa">f</span><span class="s2">&quot;Skipping normalization because distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span>
</span><span id="normalize-48"><a href="#normalize-48"><span class="linenos">48</span></a> <span class="p">)</span>
</span><span id="normalize-49"><a href="#normalize-49"><span class="linenos">49</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="normalize-50"><a href="#normalize-50"><span class="linenos">50</span></a>
</span><span id="normalize-51"><a href="#normalize-51"><span class="linenos">51</span></a> <span class="k">try</span><span class="p">:</span>
</span><span id="normalize-52"><a href="#normalize-52"><span class="linenos">52</span></a> <span class="n">node</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span>
</span><span id="normalize-53"><a href="#normalize-53"><span class="linenos">53</span></a> <span class="n">while_changing</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="normalize-54"><a href="#normalize-54"><span class="linenos">54</span></a> <span class="p">)</span>
</span><span id="normalize-55"><a href="#normalize-55"><span class="linenos">55</span></a> <span class="k">except</span> <span class="n">OptimizeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
</span><span id="normalize-56"><a href="#normalize-56"><span class="linenos">56</span></a> <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
</span><span id="normalize-57"><a href="#normalize-57"><span class="linenos">57</span></a> <span class="n">node</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">original</span><span class="p">)</span>
</span><span id="normalize-58"><a href="#normalize-58"><span class="linenos">58</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="normalize-59"><a href="#normalize-59"><span class="linenos">59</span></a> <span class="k">return</span> <span class="n">original</span>
</span><span id="normalize-60"><a href="#normalize-60"><span class="linenos">60</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="normalize-61"><a href="#normalize-61"><span class="linenos">61</span></a>
</span><span id="normalize-62"><a href="#normalize-62"><span class="linenos">62</span></a> <span class="k">if</span> <span class="n">root</span><span class="p">:</span>
</span><span id="normalize-63"><a href="#normalize-63"><span class="linenos">63</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">node</span>
</span><span id="normalize-64"><a href="#normalize-64"><span class="linenos">64</span></a>
</span><span id="normalize-65"><a href="#normalize-65"><span class="linenos">65</span></a> <span class="k">return</span> <span class="n">expression</span>
</span></pre></div>
@ -350,19 +371,62 @@
<div class="attr function">
<span class="def">def</span>
<span class="name">normalized</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">expression</span>, </span><span class="param"><span class="n">dnf</span><span class="o">=</span><span class="kc">False</span></span><span class="return-annotation">):</span></span>
<span class="name">normalized</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">expression</span><span class="p">:</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span>, </span><span class="param"><span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span></span><span class="return-annotation">) -> <span class="nb">bool</span>:</span></span>
<label class="view-source-button" for="normalized-view-source"><span>View Source</span></label>
</div>
<a class="headerlink" href="#normalized"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalized-67"><a href="#normalized-67"><span class="linenos">67</span></a><span class="k">def</span> <span class="nf">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
</span><span id="normalized-68"><a href="#normalized-68"><span class="linenos">68</span></a> <span class="n">ancestor</span><span class="p">,</span> <span class="n">root</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span>
</span><span id="normalized-69"><a href="#normalized-69"><span class="linenos">69</span></a>
</span><span id="normalized-70"><a href="#normalized-70"><span class="linenos">70</span></a> <span class="k">return</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">connector</span><span class="o">.</span><span class="n">find_ancestor</span><span class="p">(</span><span class="n">ancestor</span><span class="p">)</span> <span class="k">for</span> <span class="n">connector</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">root</span><span class="p">))</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalized-68"><a href="#normalized-68"><span class="linenos">68</span></a><span class="k">def</span> <span class="nf">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
</span><span id="normalized-69"><a href="#normalized-69"><span class="linenos">69</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalized-70"><a href="#normalized-70"><span class="linenos">70</span></a><span class="sd"> Checks whether a given expression is in a normal form of interest.</span>
</span><span id="normalized-71"><a href="#normalized-71"><span class="linenos">71</span></a>
</span><span id="normalized-72"><a href="#normalized-72"><span class="linenos">72</span></a><span class="sd"> Example:</span>
</span><span id="normalized-73"><a href="#normalized-73"><span class="linenos">73</span></a><span class="sd"> &gt;&gt;&gt; from sqlglot import parse_one</span>
</span><span id="normalized-74"><a href="#normalized-74"><span class="linenos">74</span></a><span class="sd"> &gt;&gt;&gt; normalized(parse_one(&quot;(a AND b) OR c OR (d AND e)&quot;), dnf=True)</span>
</span><span id="normalized-75"><a href="#normalized-75"><span class="linenos">75</span></a><span class="sd"> True</span>
</span><span id="normalized-76"><a href="#normalized-76"><span class="linenos">76</span></a><span class="sd"> &gt;&gt;&gt; normalized(parse_one(&quot;(a OR b) AND c&quot;)) # Checks CNF by default</span>
</span><span id="normalized-77"><a href="#normalized-77"><span class="linenos">77</span></a><span class="sd"> True</span>
</span><span id="normalized-78"><a href="#normalized-78"><span class="linenos">78</span></a><span class="sd"> &gt;&gt;&gt; normalized(parse_one(&quot;a AND (b OR c)&quot;), dnf=True)</span>
</span><span id="normalized-79"><a href="#normalized-79"><span class="linenos">79</span></a><span class="sd"> False</span>
</span><span id="normalized-80"><a href="#normalized-80"><span class="linenos">80</span></a>
</span><span id="normalized-81"><a href="#normalized-81"><span class="linenos">81</span></a><span class="sd"> Args:</span>
</span><span id="normalized-82"><a href="#normalized-82"><span class="linenos">82</span></a><span class="sd"> expression: The expression to check if it&#39;s normalized.</span>
</span><span id="normalized-83"><a href="#normalized-83"><span class="linenos">83</span></a><span class="sd"> dnf: Whether or not to check if the expression is in Disjunctive Normal Form (DNF).</span>
</span><span id="normalized-84"><a href="#normalized-84"><span class="linenos">84</span></a><span class="sd"> Default: False, i.e. we check if it&#39;s in Conjunctive Normal Form (CNF).</span>
</span><span id="normalized-85"><a href="#normalized-85"><span class="linenos">85</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalized-86"><a href="#normalized-86"><span class="linenos">86</span></a> <span class="n">ancestor</span><span class="p">,</span> <span class="n">root</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span>
</span><span id="normalized-87"><a href="#normalized-87"><span class="linenos">87</span></a> <span class="k">return</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span>
</span><span id="normalized-88"><a href="#normalized-88"><span class="linenos">88</span></a> <span class="n">connector</span><span class="o">.</span><span class="n">find_ancestor</span><span class="p">(</span><span class="n">ancestor</span><span class="p">)</span> <span class="k">for</span> <span class="n">connector</span> <span class="ow">in</span> <span class="n">find_all_in_scope</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">root</span><span class="p">)</span>
</span><span id="normalized-89"><a href="#normalized-89"><span class="linenos">89</span></a> <span class="p">)</span>
</span></pre></div>
<div class="docstring"><p>Checks whether a given expression is in a normal form of interest.</p>
<h6 id="example">Example:</h6>
<blockquote>
<div class="pdoc-code codehilite">
<pre><span></span><code><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">sqlglot</span> <span class="kn">import</span> <span class="n">parse_one</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">normalized</span><span class="p">(</span><span class="n">parse_one</span><span class="p">(</span><span class="s2">&quot;(a AND b) OR c OR (d AND e)&quot;</span><span class="p">),</span> <span class="n">dnf</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="go">True</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">normalized</span><span class="p">(</span><span class="n">parse_one</span><span class="p">(</span><span class="s2">&quot;(a OR b) AND c&quot;</span><span class="p">))</span> <span class="c1"># Checks CNF by default</span>
<span class="go">True</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">normalized</span><span class="p">(</span><span class="n">parse_one</span><span class="p">(</span><span class="s2">&quot;a AND (b OR c)&quot;</span><span class="p">),</span> <span class="n">dnf</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="go">False</span>
</code></pre>
</div>
</blockquote>
<h6 id="arguments">Arguments:</h6>
<ul>
<li><strong>expression:</strong> The expression to check if it's normalized.</li>
<li><strong>dnf:</strong> Whether or not 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).</li>
</ul>
</div>
</section>
@ -371,37 +435,39 @@
<div class="attr function">
<span class="def">def</span>
<span class="name">normalization_distance</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">expression</span>, </span><span class="param"><span class="n">dnf</span><span class="o">=</span><span class="kc">False</span></span><span class="return-annotation">):</span></span>
<span class="name">normalization_distance</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">expression</span><span class="p">:</span> <span class="n"><a href="../expressions.html#Expression">sqlglot.expressions.Expression</a></span>, </span><span class="param"><span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span></span><span class="return-annotation">) -> <span class="nb">int</span>:</span></span>
<label class="view-source-button" for="normalization_distance-view-source"><span>View Source</span></label>
</div>
<a class="headerlink" href="#normalization_distance"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalization_distance-73"><a href="#normalization_distance-73"><span class="linenos">73</span></a><span class="k">def</span> <span class="nf">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
</span><span id="normalization_distance-74"><a href="#normalization_distance-74"><span class="linenos">74</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalization_distance-75"><a href="#normalization_distance-75"><span class="linenos">75</span></a><span class="sd"> The difference in the number of predicates between the current expression and the normalized form.</span>
</span><span id="normalization_distance-76"><a href="#normalization_distance-76"><span class="linenos">76</span></a>
</span><span id="normalization_distance-77"><a href="#normalization_distance-77"><span class="linenos">77</span></a><span class="sd"> This is used as an estimate of the cost of the conversion which is exponential in complexity.</span>
</span><span id="normalization_distance-78"><a href="#normalization_distance-78"><span class="linenos">78</span></a>
</span><span id="normalization_distance-79"><a href="#normalization_distance-79"><span class="linenos">79</span></a><span class="sd"> Example:</span>
</span><span id="normalization_distance-80"><a href="#normalization_distance-80"><span class="linenos">80</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="normalization_distance-81"><a href="#normalization_distance-81"><span class="linenos">81</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(a AND b) OR (c AND d)&quot;)</span>
</span><span id="normalization_distance-82"><a href="#normalization_distance-82"><span class="linenos">82</span></a><span class="sd"> &gt;&gt;&gt; normalization_distance(expression)</span>
</span><span id="normalization_distance-83"><a href="#normalization_distance-83"><span class="linenos">83</span></a><span class="sd"> 4</span>
</span><span id="normalization_distance-84"><a href="#normalization_distance-84"><span class="linenos">84</span></a>
</span><span id="normalization_distance-85"><a href="#normalization_distance-85"><span class="linenos">85</span></a><span class="sd"> Args:</span>
</span><span id="normalization_distance-86"><a href="#normalization_distance-86"><span class="linenos">86</span></a><span class="sd"> expression (sqlglot.Expression): expression to compute distance</span>
</span><span id="normalization_distance-87"><a href="#normalization_distance-87"><span class="linenos">87</span></a><span class="sd"> dnf (bool): compute to dnf distance instead</span>
</span><span id="normalization_distance-88"><a href="#normalization_distance-88"><span class="linenos">88</span></a><span class="sd"> Returns:</span>
</span><span id="normalization_distance-89"><a href="#normalization_distance-89"><span class="linenos">89</span></a><span class="sd"> int: difference</span>
</span><span id="normalization_distance-90"><a href="#normalization_distance-90"><span class="linenos">90</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalization_distance-91"><a href="#normalization_distance-91"><span class="linenos">91</span></a> <span class="k">return</span> <span class="nb">sum</span><span class="p">(</span><span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">))</span> <span class="o">-</span> <span class="p">(</span>
</span><span id="normalization_distance-92"><a href="#normalization_distance-92"><span class="linenos">92</span></a> <span class="nb">sum</span><span class="p">(</span><span class="mi">1</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))</span> <span class="o">+</span> <span class="mi">1</span>
</span><span id="normalization_distance-93"><a href="#normalization_distance-93"><span class="linenos">93</span></a> <span class="p">)</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalization_distance-92"><a href="#normalization_distance-92"><span class="linenos"> 92</span></a><span class="k">def</span> <span class="nf">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
</span><span id="normalization_distance-93"><a href="#normalization_distance-93"><span class="linenos"> 93</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalization_distance-94"><a href="#normalization_distance-94"><span class="linenos"> 94</span></a><span class="sd"> The difference in the number of predicates between a given expression and its normalized form.</span>
</span><span id="normalization_distance-95"><a href="#normalization_distance-95"><span class="linenos"> 95</span></a>
</span><span id="normalization_distance-96"><a href="#normalization_distance-96"><span class="linenos"> 96</span></a><span class="sd"> This is used as an estimate of the cost of the conversion which is exponential in complexity.</span>
</span><span id="normalization_distance-97"><a href="#normalization_distance-97"><span class="linenos"> 97</span></a>
</span><span id="normalization_distance-98"><a href="#normalization_distance-98"><span class="linenos"> 98</span></a><span class="sd"> Example:</span>
</span><span id="normalization_distance-99"><a href="#normalization_distance-99"><span class="linenos"> 99</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="normalization_distance-100"><a href="#normalization_distance-100"><span class="linenos">100</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(&quot;(a AND b) OR (c AND d)&quot;)</span>
</span><span id="normalization_distance-101"><a href="#normalization_distance-101"><span class="linenos">101</span></a><span class="sd"> &gt;&gt;&gt; normalization_distance(expression)</span>
</span><span id="normalization_distance-102"><a href="#normalization_distance-102"><span class="linenos">102</span></a><span class="sd"> 4</span>
</span><span id="normalization_distance-103"><a href="#normalization_distance-103"><span class="linenos">103</span></a>
</span><span id="normalization_distance-104"><a href="#normalization_distance-104"><span class="linenos">104</span></a><span class="sd"> Args:</span>
</span><span id="normalization_distance-105"><a href="#normalization_distance-105"><span class="linenos">105</span></a><span class="sd"> expression: The expression to compute the normalization distance for.</span>
</span><span id="normalization_distance-106"><a href="#normalization_distance-106"><span class="linenos">106</span></a><span class="sd"> dnf: Whether or not to check if the expression is in Disjunctive Normal Form (DNF).</span>
</span><span id="normalization_distance-107"><a href="#normalization_distance-107"><span class="linenos">107</span></a><span class="sd"> Default: False, i.e. we check if it&#39;s in Conjunctive Normal Form (CNF).</span>
</span><span id="normalization_distance-108"><a href="#normalization_distance-108"><span class="linenos">108</span></a>
</span><span id="normalization_distance-109"><a href="#normalization_distance-109"><span class="linenos">109</span></a><span class="sd"> Returns:</span>
</span><span id="normalization_distance-110"><a href="#normalization_distance-110"><span class="linenos">110</span></a><span class="sd"> The normalization distance.</span>
</span><span id="normalization_distance-111"><a href="#normalization_distance-111"><span class="linenos">111</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalization_distance-112"><a href="#normalization_distance-112"><span class="linenos">112</span></a> <span class="k">return</span> <span class="nb">sum</span><span class="p">(</span><span class="n">_predicate_lengths</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">))</span> <span class="o">-</span> <span class="p">(</span>
</span><span id="normalization_distance-113"><a href="#normalization_distance-113"><span class="linenos">113</span></a> <span class="nb">sum</span><span class="p">(</span><span class="mi">1</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))</span> <span class="o">+</span> <span class="mi">1</span>
</span><span id="normalization_distance-114"><a href="#normalization_distance-114"><span class="linenos">114</span></a> <span class="p">)</span>
</span></pre></div>
<div class="docstring"><p>The difference in the number of predicates between the current expression and the normalized form.</p>
<div class="docstring"><p>The difference in the number of predicates between a given expression and its normalized form.</p>
<p>This is used as an estimate of the cost of the conversion which is exponential in complexity.</p>
@ -420,14 +486,15 @@
<h6 id="arguments">Arguments:</h6>
<ul>
<li><strong>expression (sqlglot.Expression):</strong> expression to compute distance</li>
<li><strong>dnf (bool):</strong> compute to dnf distance instead</li>
<li><strong>expression:</strong> The expression to compute the normalization distance for.</li>
<li><strong>dnf:</strong> Whether or not 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).</li>
</ul>
<h6 id="returns">Returns:</h6>
<blockquote>
<p>int: difference</p>
<p>The normalization distance.</p>
</blockquote>
</div>
@ -444,38 +511,38 @@
</div>
<a class="headerlink" href="#distributive_law"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="distributive_law-116"><a href="#distributive_law-116"><span class="linenos">116</span></a><span class="k">def</span> <span class="nf">distributive_law</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">):</span>
</span><span id="distributive_law-117"><a href="#distributive_law-117"><span class="linenos">117</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="distributive_law-118"><a href="#distributive_law-118"><span class="linenos">118</span></a><span class="sd"> x OR (y AND z) -&gt; (x OR y) AND (x OR z)</span>
</span><span id="distributive_law-119"><a href="#distributive_law-119"><span class="linenos">119</span></a><span class="sd"> (x AND y) OR (y AND z) -&gt; (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)</span>
</span><span id="distributive_law-120"><a href="#distributive_law-120"><span class="linenos">120</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="distributive_law-121"><a href="#distributive_law-121"><span class="linenos">121</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="distributive_law-122"><a href="#distributive_law-122"><span class="linenos">122</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="distributive_law-123"><a href="#distributive_law-123"><span class="linenos">123</span></a>
</span><span id="distributive_law-124"><a href="#distributive_law-124"><span class="linenos">124</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="distributive_law-125"><a href="#distributive_law-125"><span class="linenos">125</span></a>
</span><span id="distributive_law-126"><a href="#distributive_law-126"><span class="linenos">126</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="distributive_law-127"><a href="#distributive_law-127"><span class="linenos">127</span></a> <span class="k">raise</span> <span class="n">OptimizeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Normalization distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
</span><span id="distributive_law-128"><a href="#distributive_law-128"><span class="linenos">128</span></a>
</span><span id="distributive_law-129"><a href="#distributive_law-129"><span class="linenos">129</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">replace_children</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="distributive_law-130"><a href="#distributive_law-130"><span class="linenos">130</span></a> <span class="n">to_exp</span><span class="p">,</span> <span class="n">from_exp</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span>
</span><span id="distributive_law-131"><a href="#distributive_law-131"><span class="linenos">131</span></a>
</span><span id="distributive_law-132"><a href="#distributive_law-132"><span class="linenos">132</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">from_exp</span><span class="p">):</span>
</span><span id="distributive_law-133"><a href="#distributive_law-133"><span class="linenos">133</span></a> <span class="n">a</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">unnest_operands</span><span class="p">()</span>
</span><span id="distributive_law-134"><a href="#distributive_law-134"><span class="linenos">134</span></a>
</span><span id="distributive_law-135"><a href="#distributive_law-135"><span class="linenos">135</span></a> <span class="n">from_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">from_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="distributive_law-136"><a href="#distributive_law-136"><span class="linenos">136</span></a> <span class="n">to_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">to_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="distributive_law-137"><a href="#distributive_law-137"><span class="linenos">137</span></a>
</span><span id="distributive_law-138"><a href="#distributive_law-138"><span class="linenos">138</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="distributive_law-139"><a href="#distributive_law-139"><span class="linenos">139</span></a> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">)))</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">b</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="distributive_law-140"><a href="#distributive_law-140"><span class="linenos">140</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-141"><a href="#distributive_law-141"><span class="linenos">141</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-142"><a href="#distributive_law-142"><span class="linenos">142</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="distributive_law-143"><a href="#distributive_law-143"><span class="linenos">143</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-144"><a href="#distributive_law-144"><span class="linenos">144</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="distributive_law-145"><a href="#distributive_law-145"><span class="linenos">145</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="distributive_law-137"><a href="#distributive_law-137"><span class="linenos">137</span></a><span class="k">def</span> <span class="nf">distributive_law</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">):</span>
</span><span id="distributive_law-138"><a href="#distributive_law-138"><span class="linenos">138</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="distributive_law-139"><a href="#distributive_law-139"><span class="linenos">139</span></a><span class="sd"> x OR (y AND z) -&gt; (x OR y) AND (x OR z)</span>
</span><span id="distributive_law-140"><a href="#distributive_law-140"><span class="linenos">140</span></a><span class="sd"> (x AND y) OR (y AND z) -&gt; (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)</span>
</span><span id="distributive_law-141"><a href="#distributive_law-141"><span class="linenos">141</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="distributive_law-142"><a href="#distributive_law-142"><span class="linenos">142</span></a> <span class="k">if</span> <span class="n">normalized</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">):</span>
</span><span id="distributive_law-143"><a href="#distributive_law-143"><span class="linenos">143</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="distributive_law-144"><a href="#distributive_law-144"><span class="linenos">144</span></a>
</span><span id="distributive_law-145"><a href="#distributive_law-145"><span class="linenos">145</span></a> <span class="n">distance</span> <span class="o">=</span> <span class="n">normalization_distance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dnf</span><span class="o">=</span><span class="n">dnf</span><span class="p">)</span>
</span><span id="distributive_law-146"><a href="#distributive_law-146"><span class="linenos">146</span></a>
</span><span id="distributive_law-147"><a href="#distributive_law-147"><span class="linenos">147</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="distributive_law-147"><a href="#distributive_law-147"><span class="linenos">147</span></a> <span class="k">if</span> <span class="n">distance</span> <span class="o">&gt;</span> <span class="n">max_distance</span><span class="p">:</span>
</span><span id="distributive_law-148"><a href="#distributive_law-148"><span class="linenos">148</span></a> <span class="k">raise</span> <span class="n">OptimizeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Normalization distance </span><span class="si">{</span><span class="n">distance</span><span class="si">}</span><span class="s2"> exceeds max </span><span class="si">{</span><span class="n">max_distance</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
</span><span id="distributive_law-149"><a href="#distributive_law-149"><span class="linenos">149</span></a>
</span><span id="distributive_law-150"><a href="#distributive_law-150"><span class="linenos">150</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">replace_children</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">e</span><span class="p">:</span> <span class="n">distributive_law</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">dnf</span><span class="p">,</span> <span class="n">max_distance</span><span class="p">,</span> <span class="n">generate</span><span class="p">))</span>
</span><span id="distributive_law-151"><a href="#distributive_law-151"><span class="linenos">151</span></a> <span class="n">to_exp</span><span class="p">,</span> <span class="n">from_exp</span> <span class="o">=</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">)</span> <span class="k">if</span> <span class="n">dnf</span> <span class="k">else</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">And</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Or</span><span class="p">)</span>
</span><span id="distributive_law-152"><a href="#distributive_law-152"><span class="linenos">152</span></a>
</span><span id="distributive_law-153"><a href="#distributive_law-153"><span class="linenos">153</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">from_exp</span><span class="p">):</span>
</span><span id="distributive_law-154"><a href="#distributive_law-154"><span class="linenos">154</span></a> <span class="n">a</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">expression</span><span class="o">.</span><span class="n">unnest_operands</span><span class="p">()</span>
</span><span id="distributive_law-155"><a href="#distributive_law-155"><span class="linenos">155</span></a>
</span><span id="distributive_law-156"><a href="#distributive_law-156"><span class="linenos">156</span></a> <span class="n">from_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">from_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="distributive_law-157"><a href="#distributive_law-157"><span class="linenos">157</span></a> <span class="n">to_func</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">and_</span> <span class="k">if</span> <span class="n">to_exp</span> <span class="o">==</span> <span class="n">exp</span><span class="o">.</span><span class="n">And</span> <span class="k">else</span> <span class="n">exp</span><span class="o">.</span><span class="n">or_</span>
</span><span id="distributive_law-158"><a href="#distributive_law-158"><span class="linenos">158</span></a>
</span><span id="distributive_law-159"><a href="#distributive_law-159"><span class="linenos">159</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="distributive_law-160"><a href="#distributive_law-160"><span class="linenos">160</span></a> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">)))</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">b</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Connector</span><span class="p">))):</span>
</span><span id="distributive_law-161"><a href="#distributive_law-161"><span class="linenos">161</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-162"><a href="#distributive_law-162"><span class="linenos">162</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-163"><a href="#distributive_law-163"><span class="linenos">163</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="distributive_law-164"><a href="#distributive_law-164"><span class="linenos">164</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-165"><a href="#distributive_law-165"><span class="linenos">165</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">to_exp</span><span class="p">):</span>
</span><span id="distributive_law-166"><a href="#distributive_law-166"><span class="linenos">166</span></a> <span class="k">return</span> <span class="n">_distribute</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">from_func</span><span class="p">,</span> <span class="n">to_func</span><span class="p">,</span> <span class="n">generate</span><span class="p">)</span>
</span><span id="distributive_law-167"><a href="#distributive_law-167"><span class="linenos">167</span></a>
</span><span id="distributive_law-168"><a href="#distributive_law-168"><span class="linenos">168</span></a> <span class="k">return</span> <span class="n">expression</span>
</span></pre></div>

View file

@ -60,7 +60,7 @@
</span><span id="L-2"><a href="#L-2"><span class="linenos"> 2</span></a>
</span><span id="L-3"><a href="#L-3"><span class="linenos"> 3</span></a><span class="kn">import</span> <span class="nn">typing</span> <span class="k">as</span> <span class="nn">t</span>
</span><span id="L-4"><a href="#L-4"><span class="linenos"> 4</span></a>
</span><span id="L-5"><a href="#L-5"><span class="linenos"> 5</span></a><span class="kn">from</span> <span class="nn">sqlglot</span> <span class="kn">import</span> <span class="n">exp</span>
</span><span id="L-5"><a href="#L-5"><span class="linenos"> 5</span></a><span class="kn">from</span> <span class="nn">sqlglot</span> <span class="kn">import</span> <span class="n">exp</span><span class="p">,</span> <span class="n">parse_one</span>
</span><span id="L-6"><a href="#L-6"><span class="linenos"> 6</span></a><span class="kn">from</span> <span class="nn">sqlglot._typing</span> <span class="kn">import</span> <span class="n">E</span>
</span><span id="L-7"><a href="#L-7"><span class="linenos"> 7</span></a><span class="kn">from</span> <span class="nn">sqlglot.dialects.dialect</span> <span class="kn">import</span> <span class="n">Dialect</span><span class="p">,</span> <span class="n">DialectType</span>
</span><span id="L-8"><a href="#L-8"><span class="linenos"> 8</span></a>
@ -107,7 +107,7 @@
</span><span id="L-49"><a href="#L-49"><span class="linenos">49</span></a><span class="sd"> The transformed expression.</span>
</span><span id="L-50"><a href="#L-50"><span class="linenos">50</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-51"><a href="#L-51"><span class="linenos">51</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos">52</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">to_identifier</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos">52</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">parse_one</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dialect</span><span class="o">=</span><span class="n">dialect</span><span class="p">,</span> <span class="n">into</span><span class="o">=</span><span class="n">exp</span><span class="o">.</span><span class="n">Identifier</span><span class="p">)</span>
</span><span id="L-53"><a href="#L-53"><span class="linenos">53</span></a>
</span><span id="L-54"><a href="#L-54"><span class="linenos">54</span></a> <span class="n">dialect</span> <span class="o">=</span> <span class="n">Dialect</span><span class="o">.</span><span class="n">get_or_raise</span><span class="p">(</span><span class="n">dialect</span><span class="p">)</span>
</span><span id="L-55"><a href="#L-55"><span class="linenos">55</span></a>
@ -165,7 +165,7 @@
</span><span id="normalize_identifiers-50"><a href="#normalize_identifiers-50"><span class="linenos">50</span></a><span class="sd"> The transformed expression.</span>
</span><span id="normalize_identifiers-51"><a href="#normalize_identifiers-51"><span class="linenos">51</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalize_identifiers-52"><a href="#normalize_identifiers-52"><span class="linenos">52</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
</span><span id="normalize_identifiers-53"><a href="#normalize_identifiers-53"><span class="linenos">53</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">to_identifier</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="normalize_identifiers-53"><a href="#normalize_identifiers-53"><span class="linenos">53</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">parse_one</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">dialect</span><span class="o">=</span><span class="n">dialect</span><span class="p">,</span> <span class="n">into</span><span class="o">=</span><span class="n">exp</span><span class="o">.</span><span class="n">Identifier</span><span class="p">)</span>
</span><span id="normalize_identifiers-54"><a href="#normalize_identifiers-54"><span class="linenos">54</span></a>
</span><span id="normalize_identifiers-55"><a href="#normalize_identifiers-55"><span class="linenos">55</span></a> <span class="n">dialect</span> <span class="o">=</span> <span class="n">Dialect</span><span class="o">.</span><span class="n">get_or_raise</span><span class="p">(</span><span class="n">dialect</span><span class="p">)</span>
</span><span id="normalize_identifiers-56"><a href="#normalize_identifiers-56"><span class="linenos">56</span></a>

View file

@ -109,52 +109,56 @@
</span><span id="L-39"><a href="#L-39"><span class="linenos">39</span></a> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">other_table_names</span><span class="p">(</span><span class="n">dep</span><span class="p">))</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">:</span>
</span><span id="L-40"><a href="#L-40"><span class="linenos">40</span></a> <span class="k">continue</span>
</span><span id="L-41"><a href="#L-41"><span class="linenos">41</span></a>
</span><span id="L-42"><a href="#L-42"><span class="linenos">42</span></a> <span class="k">for</span> <span class="n">predicate</span> <span class="ow">in</span> <span class="n">on</span><span class="o">.</span><span class="n">flatten</span><span class="p">():</span>
</span><span id="L-43"><a href="#L-43"><span class="linenos">43</span></a> <span class="k">if</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">predicate</span><span class="p">):</span>
</span><span id="L-44"><a href="#L-44"><span class="linenos">44</span></a> <span class="n">predicate</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="L-45"><a href="#L-45"><span class="linenos">45</span></a> <span class="n">join</span><span class="o">.</span><span class="n">on</span><span class="p">(</span><span class="n">predicate</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-46"><a href="#L-46"><span class="linenos">46</span></a>
</span><span id="L-47"><a href="#L-47"><span class="linenos">47</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="L-48"><a href="#L-48"><span class="linenos">48</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="L-49"><a href="#L-49"><span class="linenos">49</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-42"><a href="#L-42"><span class="linenos">42</span></a> <span class="n">operator</span> <span class="o">=</span> <span class="nb">type</span><span class="p">(</span><span class="n">on</span><span class="p">)</span>
</span><span id="L-43"><a href="#L-43"><span class="linenos">43</span></a> <span class="k">for</span> <span class="n">predicate</span> <span class="ow">in</span> <span class="n">on</span><span class="o">.</span><span class="n">flatten</span><span class="p">():</span>
</span><span id="L-44"><a href="#L-44"><span class="linenos">44</span></a> <span class="k">if</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">predicate</span><span class="p">):</span>
</span><span id="L-45"><a href="#L-45"><span class="linenos">45</span></a> <span class="n">predicate</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="L-46"><a href="#L-46"><span class="linenos">46</span></a> <span class="n">predicate</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">_combine</span><span class="p">(</span>
</span><span id="L-47"><a href="#L-47"><span class="linenos">47</span></a> <span class="p">[</span><span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">),</span> <span class="n">predicate</span><span class="p">],</span> <span class="n">operator</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span>
</span><span id="L-48"><a href="#L-48"><span class="linenos">48</span></a> <span class="p">)</span>
</span><span id="L-49"><a href="#L-49"><span class="linenos">49</span></a> <span class="n">join</span><span class="o">.</span><span class="n">on</span><span class="p">(</span><span class="n">predicate</span><span class="p">,</span> <span class="n">append</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-50"><a href="#L-50"><span class="linenos">50</span></a>
</span><span id="L-51"><a href="#L-51"><span class="linenos">51</span></a>
</span><span id="L-52"><a href="#L-52"><span class="linenos">52</span></a><span class="k">def</span> <span class="nf">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="L-53"><a href="#L-53"><span class="linenos">53</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-54"><a href="#L-54"><span class="linenos">54</span></a><span class="sd"> Reorder joins by topological sort order based on predicate references.</span>
</span><span id="L-55"><a href="#L-55"><span class="linenos">55</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-56"><a href="#L-56"><span class="linenos">56</span></a> <span class="k">for</span> <span class="n">from_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">From</span><span class="p">):</span>
</span><span id="L-57"><a href="#L-57"><span class="linenos">57</span></a> <span class="n">parent</span> <span class="o">=</span> <span class="n">from_</span><span class="o">.</span><span class="n">parent</span>
</span><span id="L-58"><a href="#L-58"><span class="linenos">58</span></a> <span class="n">joins</span> <span class="o">=</span> <span class="p">{</span><span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">:</span> <span class="n">join</span> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">parent</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;joins&quot;</span><span class="p">,</span> <span class="p">[])}</span>
</span><span id="L-59"><a href="#L-59"><span class="linenos">59</span></a> <span class="n">dag</span> <span class="o">=</span> <span class="p">{</span><span class="n">name</span><span class="p">:</span> <span class="n">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">joins</span><span class="o">.</span><span class="n">items</span><span class="p">()}</span>
</span><span id="L-60"><a href="#L-60"><span class="linenos">60</span></a> <span class="n">parent</span><span class="o">.</span><span class="n">set</span><span class="p">(</span>
</span><span id="L-61"><a href="#L-61"><span class="linenos">61</span></a> <span class="s2">&quot;joins&quot;</span><span class="p">,</span>
</span><span id="L-62"><a href="#L-62"><span class="linenos">62</span></a> <span class="p">[</span><span class="n">joins</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">tsort</span><span class="p">(</span><span class="n">dag</span><span class="p">)</span> <span class="k">if</span> <span class="n">name</span> <span class="o">!=</span> <span class="n">from_</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">and</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">joins</span><span class="p">],</span>
</span><span id="L-63"><a href="#L-63"><span class="linenos">63</span></a> <span class="p">)</span>
</span><span id="L-64"><a href="#L-64"><span class="linenos">64</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-65"><a href="#L-65"><span class="linenos">65</span></a>
</span><span id="L-66"><a href="#L-66"><span class="linenos">66</span></a>
</span><span id="L-67"><a href="#L-67"><span class="linenos">67</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="L-68"><a href="#L-68"><span class="linenos">68</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-69"><a href="#L-69"><span class="linenos">69</span></a><span class="sd"> Remove INNER and OUTER from joins as they are optional.</span>
</span><span id="L-70"><a href="#L-70"><span class="linenos">70</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-71"><a href="#L-71"><span class="linenos">71</span></a> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">):</span>
</span><span id="L-72"><a href="#L-72"><span class="linenos">72</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">k</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">JOIN_ATTRS</span><span class="p">):</span>
</span><span id="L-73"><a href="#L-73"><span class="linenos">73</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">)</span>
</span><span id="L-74"><a href="#L-74"><span class="linenos">74</span></a>
</span><span id="L-75"><a href="#L-75"><span class="linenos">75</span></a> <span class="k">if</span> <span class="n">join</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">:</span>
</span><span id="L-76"><a href="#L-76"><span class="linenos">76</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="L-77"><a href="#L-77"><span class="linenos">77</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-78"><a href="#L-78"><span class="linenos">78</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="L-79"><a href="#L-79"><span class="linenos">79</span></a>
</span><span id="L-80"><a href="#L-80"><span class="linenos">80</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;using&quot;</span><span class="p">):</span>
</span><span id="L-81"><a href="#L-81"><span class="linenos">81</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="L-82"><a href="#L-82"><span class="linenos">82</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-51"><a href="#L-51"><span class="linenos">51</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos">52</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="L-53"><a href="#L-53"><span class="linenos">53</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-54"><a href="#L-54"><span class="linenos">54</span></a>
</span><span id="L-55"><a href="#L-55"><span class="linenos">55</span></a>
</span><span id="L-56"><a href="#L-56"><span class="linenos">56</span></a><span class="k">def</span> <span class="nf">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="L-57"><a href="#L-57"><span class="linenos">57</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-58"><a href="#L-58"><span class="linenos">58</span></a><span class="sd"> Reorder joins by topological sort order based on predicate references.</span>
</span><span id="L-59"><a href="#L-59"><span class="linenos">59</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-60"><a href="#L-60"><span class="linenos">60</span></a> <span class="k">for</span> <span class="n">from_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">From</span><span class="p">):</span>
</span><span id="L-61"><a href="#L-61"><span class="linenos">61</span></a> <span class="n">parent</span> <span class="o">=</span> <span class="n">from_</span><span class="o">.</span><span class="n">parent</span>
</span><span id="L-62"><a href="#L-62"><span class="linenos">62</span></a> <span class="n">joins</span> <span class="o">=</span> <span class="p">{</span><span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">:</span> <span class="n">join</span> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">parent</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;joins&quot;</span><span class="p">,</span> <span class="p">[])}</span>
</span><span id="L-63"><a href="#L-63"><span class="linenos">63</span></a> <span class="n">dag</span> <span class="o">=</span> <span class="p">{</span><span class="n">name</span><span class="p">:</span> <span class="n">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">joins</span><span class="o">.</span><span class="n">items</span><span class="p">()}</span>
</span><span id="L-64"><a href="#L-64"><span class="linenos">64</span></a> <span class="n">parent</span><span class="o">.</span><span class="n">set</span><span class="p">(</span>
</span><span id="L-65"><a href="#L-65"><span class="linenos">65</span></a> <span class="s2">&quot;joins&quot;</span><span class="p">,</span>
</span><span id="L-66"><a href="#L-66"><span class="linenos">66</span></a> <span class="p">[</span><span class="n">joins</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">tsort</span><span class="p">(</span><span class="n">dag</span><span class="p">)</span> <span class="k">if</span> <span class="n">name</span> <span class="o">!=</span> <span class="n">from_</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">and</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">joins</span><span class="p">],</span>
</span><span id="L-67"><a href="#L-67"><span class="linenos">67</span></a> <span class="p">)</span>
</span><span id="L-68"><a href="#L-68"><span class="linenos">68</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-69"><a href="#L-69"><span class="linenos">69</span></a>
</span><span id="L-70"><a href="#L-70"><span class="linenos">70</span></a>
</span><span id="L-71"><a href="#L-71"><span class="linenos">71</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="L-72"><a href="#L-72"><span class="linenos">72</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-73"><a href="#L-73"><span class="linenos">73</span></a><span class="sd"> Remove INNER and OUTER from joins as they are optional.</span>
</span><span id="L-74"><a href="#L-74"><span class="linenos">74</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-75"><a href="#L-75"><span class="linenos">75</span></a> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">):</span>
</span><span id="L-76"><a href="#L-76"><span class="linenos">76</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">k</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">JOIN_ATTRS</span><span class="p">):</span>
</span><span id="L-77"><a href="#L-77"><span class="linenos">77</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">)</span>
</span><span id="L-78"><a href="#L-78"><span class="linenos">78</span></a>
</span><span id="L-79"><a href="#L-79"><span class="linenos">79</span></a> <span class="k">if</span> <span class="n">join</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">:</span>
</span><span id="L-80"><a href="#L-80"><span class="linenos">80</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="L-81"><a href="#L-81"><span class="linenos">81</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-82"><a href="#L-82"><span class="linenos">82</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="L-83"><a href="#L-83"><span class="linenos">83</span></a>
</span><span id="L-84"><a href="#L-84"><span class="linenos">84</span></a>
</span><span id="L-85"><a href="#L-85"><span class="linenos">85</span></a><span class="k">def</span> <span class="nf">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">t</span><span class="o">.</span><span class="n">Set</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span>
</span><span id="L-86"><a href="#L-86"><span class="linenos">86</span></a> <span class="n">on</span> <span class="o">=</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span>
</span><span id="L-87"><a href="#L-87"><span class="linenos">87</span></a> <span class="k">return</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">on</span><span class="p">,</span> <span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">)</span> <span class="k">if</span> <span class="n">on</span> <span class="k">else</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="L-84"><a href="#L-84"><span class="linenos">84</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;using&quot;</span><span class="p">):</span>
</span><span id="L-85"><a href="#L-85"><span class="linenos">85</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="L-86"><a href="#L-86"><span class="linenos">86</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-87"><a href="#L-87"><span class="linenos">87</span></a>
</span><span id="L-88"><a href="#L-88"><span class="linenos">88</span></a>
</span><span id="L-89"><a href="#L-89"><span class="linenos">89</span></a><span class="k">def</span> <span class="nf">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">t</span><span class="o">.</span><span class="n">Set</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span>
</span><span id="L-90"><a href="#L-90"><span class="linenos">90</span></a> <span class="n">on</span> <span class="o">=</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span>
</span><span id="L-91"><a href="#L-91"><span class="linenos">91</span></a> <span class="k">return</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">on</span><span class="p">,</span> <span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">)</span> <span class="k">if</span> <span class="n">on</span> <span class="k">else</span> <span class="nb">set</span><span class="p">()</span>
</span></pre></div>
@ -213,14 +217,18 @@
</span><span id="optimize_joins-40"><a href="#optimize_joins-40"><span class="linenos">40</span></a> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">other_table_names</span><span class="p">(</span><span class="n">dep</span><span class="p">))</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">:</span>
</span><span id="optimize_joins-41"><a href="#optimize_joins-41"><span class="linenos">41</span></a> <span class="k">continue</span>
</span><span id="optimize_joins-42"><a href="#optimize_joins-42"><span class="linenos">42</span></a>
</span><span id="optimize_joins-43"><a href="#optimize_joins-43"><span class="linenos">43</span></a> <span class="k">for</span> <span class="n">predicate</span> <span class="ow">in</span> <span class="n">on</span><span class="o">.</span><span class="n">flatten</span><span class="p">():</span>
</span><span id="optimize_joins-44"><a href="#optimize_joins-44"><span class="linenos">44</span></a> <span class="k">if</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">predicate</span><span class="p">):</span>
</span><span id="optimize_joins-45"><a href="#optimize_joins-45"><span class="linenos">45</span></a> <span class="n">predicate</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="optimize_joins-46"><a href="#optimize_joins-46"><span class="linenos">46</span></a> <span class="n">join</span><span class="o">.</span><span class="n">on</span><span class="p">(</span><span class="n">predicate</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="optimize_joins-47"><a href="#optimize_joins-47"><span class="linenos">47</span></a>
</span><span id="optimize_joins-48"><a href="#optimize_joins-48"><span class="linenos">48</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="optimize_joins-49"><a href="#optimize_joins-49"><span class="linenos">49</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="optimize_joins-50"><a href="#optimize_joins-50"><span class="linenos">50</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="optimize_joins-43"><a href="#optimize_joins-43"><span class="linenos">43</span></a> <span class="n">operator</span> <span class="o">=</span> <span class="nb">type</span><span class="p">(</span><span class="n">on</span><span class="p">)</span>
</span><span id="optimize_joins-44"><a href="#optimize_joins-44"><span class="linenos">44</span></a> <span class="k">for</span> <span class="n">predicate</span> <span class="ow">in</span> <span class="n">on</span><span class="o">.</span><span class="n">flatten</span><span class="p">():</span>
</span><span id="optimize_joins-45"><a href="#optimize_joins-45"><span class="linenos">45</span></a> <span class="k">if</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">predicate</span><span class="p">):</span>
</span><span id="optimize_joins-46"><a href="#optimize_joins-46"><span class="linenos">46</span></a> <span class="n">predicate</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="optimize_joins-47"><a href="#optimize_joins-47"><span class="linenos">47</span></a> <span class="n">predicate</span> <span class="o">=</span> <span class="n">exp</span><span class="o">.</span><span class="n">_combine</span><span class="p">(</span>
</span><span id="optimize_joins-48"><a href="#optimize_joins-48"><span class="linenos">48</span></a> <span class="p">[</span><span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">),</span> <span class="n">predicate</span><span class="p">],</span> <span class="n">operator</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span>
</span><span id="optimize_joins-49"><a href="#optimize_joins-49"><span class="linenos">49</span></a> <span class="p">)</span>
</span><span id="optimize_joins-50"><a href="#optimize_joins-50"><span class="linenos">50</span></a> <span class="n">join</span><span class="o">.</span><span class="n">on</span><span class="p">(</span><span class="n">predicate</span><span class="p">,</span> <span class="n">append</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="optimize_joins-51"><a href="#optimize_joins-51"><span class="linenos">51</span></a>
</span><span id="optimize_joins-52"><a href="#optimize_joins-52"><span class="linenos">52</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="optimize_joins-53"><a href="#optimize_joins-53"><span class="linenos">53</span></a> <span class="n">expression</span> <span class="o">=</span> <span class="n">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">)</span>
</span><span id="optimize_joins-54"><a href="#optimize_joins-54"><span class="linenos">54</span></a> <span class="k">return</span> <span class="n">expression</span>
</span></pre></div>
@ -251,19 +259,19 @@
</div>
<a class="headerlink" href="#reorder_joins"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="reorder_joins-53"><a href="#reorder_joins-53"><span class="linenos">53</span></a><span class="k">def</span> <span class="nf">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="reorder_joins-54"><a href="#reorder_joins-54"><span class="linenos">54</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="reorder_joins-55"><a href="#reorder_joins-55"><span class="linenos">55</span></a><span class="sd"> Reorder joins by topological sort order based on predicate references.</span>
</span><span id="reorder_joins-56"><a href="#reorder_joins-56"><span class="linenos">56</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="reorder_joins-57"><a href="#reorder_joins-57"><span class="linenos">57</span></a> <span class="k">for</span> <span class="n">from_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">From</span><span class="p">):</span>
</span><span id="reorder_joins-58"><a href="#reorder_joins-58"><span class="linenos">58</span></a> <span class="n">parent</span> <span class="o">=</span> <span class="n">from_</span><span class="o">.</span><span class="n">parent</span>
</span><span id="reorder_joins-59"><a href="#reorder_joins-59"><span class="linenos">59</span></a> <span class="n">joins</span> <span class="o">=</span> <span class="p">{</span><span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">:</span> <span class="n">join</span> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">parent</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;joins&quot;</span><span class="p">,</span> <span class="p">[])}</span>
</span><span id="reorder_joins-60"><a href="#reorder_joins-60"><span class="linenos">60</span></a> <span class="n">dag</span> <span class="o">=</span> <span class="p">{</span><span class="n">name</span><span class="p">:</span> <span class="n">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">joins</span><span class="o">.</span><span class="n">items</span><span class="p">()}</span>
</span><span id="reorder_joins-61"><a href="#reorder_joins-61"><span class="linenos">61</span></a> <span class="n">parent</span><span class="o">.</span><span class="n">set</span><span class="p">(</span>
</span><span id="reorder_joins-62"><a href="#reorder_joins-62"><span class="linenos">62</span></a> <span class="s2">&quot;joins&quot;</span><span class="p">,</span>
</span><span id="reorder_joins-63"><a href="#reorder_joins-63"><span class="linenos">63</span></a> <span class="p">[</span><span class="n">joins</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">tsort</span><span class="p">(</span><span class="n">dag</span><span class="p">)</span> <span class="k">if</span> <span class="n">name</span> <span class="o">!=</span> <span class="n">from_</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">and</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">joins</span><span class="p">],</span>
</span><span id="reorder_joins-64"><a href="#reorder_joins-64"><span class="linenos">64</span></a> <span class="p">)</span>
</span><span id="reorder_joins-65"><a href="#reorder_joins-65"><span class="linenos">65</span></a> <span class="k">return</span> <span class="n">expression</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="reorder_joins-57"><a href="#reorder_joins-57"><span class="linenos">57</span></a><span class="k">def</span> <span class="nf">reorder_joins</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="reorder_joins-58"><a href="#reorder_joins-58"><span class="linenos">58</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="reorder_joins-59"><a href="#reorder_joins-59"><span class="linenos">59</span></a><span class="sd"> Reorder joins by topological sort order based on predicate references.</span>
</span><span id="reorder_joins-60"><a href="#reorder_joins-60"><span class="linenos">60</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="reorder_joins-61"><a href="#reorder_joins-61"><span class="linenos">61</span></a> <span class="k">for</span> <span class="n">from_</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">From</span><span class="p">):</span>
</span><span id="reorder_joins-62"><a href="#reorder_joins-62"><span class="linenos">62</span></a> <span class="n">parent</span> <span class="o">=</span> <span class="n">from_</span><span class="o">.</span><span class="n">parent</span>
</span><span id="reorder_joins-63"><a href="#reorder_joins-63"><span class="linenos">63</span></a> <span class="n">joins</span> <span class="o">=</span> <span class="p">{</span><span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">:</span> <span class="n">join</span> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">parent</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;joins&quot;</span><span class="p">,</span> <span class="p">[])}</span>
</span><span id="reorder_joins-64"><a href="#reorder_joins-64"><span class="linenos">64</span></a> <span class="n">dag</span> <span class="o">=</span> <span class="p">{</span><span class="n">name</span><span class="p">:</span> <span class="n">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">joins</span><span class="o">.</span><span class="n">items</span><span class="p">()}</span>
</span><span id="reorder_joins-65"><a href="#reorder_joins-65"><span class="linenos">65</span></a> <span class="n">parent</span><span class="o">.</span><span class="n">set</span><span class="p">(</span>
</span><span id="reorder_joins-66"><a href="#reorder_joins-66"><span class="linenos">66</span></a> <span class="s2">&quot;joins&quot;</span><span class="p">,</span>
</span><span id="reorder_joins-67"><a href="#reorder_joins-67"><span class="linenos">67</span></a> <span class="p">[</span><span class="n">joins</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">tsort</span><span class="p">(</span><span class="n">dag</span><span class="p">)</span> <span class="k">if</span> <span class="n">name</span> <span class="o">!=</span> <span class="n">from_</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">and</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">joins</span><span class="p">],</span>
</span><span id="reorder_joins-68"><a href="#reorder_joins-68"><span class="linenos">68</span></a> <span class="p">)</span>
</span><span id="reorder_joins-69"><a href="#reorder_joins-69"><span class="linenos">69</span></a> <span class="k">return</span> <span class="n">expression</span>
</span></pre></div>
@ -283,22 +291,22 @@
</div>
<a class="headerlink" href="#normalize"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalize-68"><a href="#normalize-68"><span class="linenos">68</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="normalize-69"><a href="#normalize-69"><span class="linenos">69</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalize-70"><a href="#normalize-70"><span class="linenos">70</span></a><span class="sd"> Remove INNER and OUTER from joins as they are optional.</span>
</span><span id="normalize-71"><a href="#normalize-71"><span class="linenos">71</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalize-72"><a href="#normalize-72"><span class="linenos">72</span></a> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">):</span>
</span><span id="normalize-73"><a href="#normalize-73"><span class="linenos">73</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">k</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">JOIN_ATTRS</span><span class="p">):</span>
</span><span id="normalize-74"><a href="#normalize-74"><span class="linenos">74</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">)</span>
</span><span id="normalize-75"><a href="#normalize-75"><span class="linenos">75</span></a>
</span><span id="normalize-76"><a href="#normalize-76"><span class="linenos">76</span></a> <span class="k">if</span> <span class="n">join</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">:</span>
</span><span id="normalize-77"><a href="#normalize-77"><span class="linenos">77</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="normalize-78"><a href="#normalize-78"><span class="linenos">78</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="normalize-79"><a href="#normalize-79"><span class="linenos">79</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="normalize-80"><a href="#normalize-80"><span class="linenos">80</span></a>
</span><span id="normalize-81"><a href="#normalize-81"><span class="linenos">81</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;using&quot;</span><span class="p">):</span>
</span><span id="normalize-82"><a href="#normalize-82"><span class="linenos">82</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="normalize-83"><a href="#normalize-83"><span class="linenos">83</span></a> <span class="k">return</span> <span class="n">expression</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="normalize-72"><a href="#normalize-72"><span class="linenos">72</span></a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">expression</span><span class="p">):</span>
</span><span id="normalize-73"><a href="#normalize-73"><span class="linenos">73</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="normalize-74"><a href="#normalize-74"><span class="linenos">74</span></a><span class="sd"> Remove INNER and OUTER from joins as they are optional.</span>
</span><span id="normalize-75"><a href="#normalize-75"><span class="linenos">75</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="normalize-76"><a href="#normalize-76"><span class="linenos">76</span></a> <span class="k">for</span> <span class="n">join</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">):</span>
</span><span id="normalize-77"><a href="#normalize-77"><span class="linenos">77</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">k</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">JOIN_ATTRS</span><span class="p">):</span>
</span><span id="normalize-78"><a href="#normalize-78"><span class="linenos">78</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">)</span>
</span><span id="normalize-79"><a href="#normalize-79"><span class="linenos">79</span></a>
</span><span id="normalize-80"><a href="#normalize-80"><span class="linenos">80</span></a> <span class="k">if</span> <span class="n">join</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="s2">&quot;CROSS&quot;</span><span class="p">:</span>
</span><span id="normalize-81"><a href="#normalize-81"><span class="linenos">81</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="normalize-82"><a href="#normalize-82"><span class="linenos">82</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="normalize-83"><a href="#normalize-83"><span class="linenos">83</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;kind&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
</span><span id="normalize-84"><a href="#normalize-84"><span class="linenos">84</span></a>
</span><span id="normalize-85"><a href="#normalize-85"><span class="linenos">85</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;using&quot;</span><span class="p">):</span>
</span><span id="normalize-86"><a href="#normalize-86"><span class="linenos">86</span></a> <span class="n">join</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">true</span><span class="p">())</span>
</span><span id="normalize-87"><a href="#normalize-87"><span class="linenos">87</span></a> <span class="k">return</span> <span class="n">expression</span>
</span></pre></div>
@ -318,9 +326,9 @@
</div>
<a class="headerlink" href="#other_table_names"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="other_table_names-86"><a href="#other_table_names-86"><span class="linenos">86</span></a><span class="k">def</span> <span class="nf">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">t</span><span class="o">.</span><span class="n">Set</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span>
</span><span id="other_table_names-87"><a href="#other_table_names-87"><span class="linenos">87</span></a> <span class="n">on</span> <span class="o">=</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span>
</span><span id="other_table_names-88"><a href="#other_table_names-88"><span class="linenos">88</span></a> <span class="k">return</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">on</span><span class="p">,</span> <span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">)</span> <span class="k">if</span> <span class="n">on</span> <span class="k">else</span> <span class="nb">set</span><span class="p">()</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="other_table_names-90"><a href="#other_table_names-90"><span class="linenos">90</span></a><span class="k">def</span> <span class="nf">other_table_names</span><span class="p">(</span><span class="n">join</span><span class="p">:</span> <span class="n">exp</span><span class="o">.</span><span class="n">Join</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">t</span><span class="o">.</span><span class="n">Set</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span>
</span><span id="other_table_names-91"><a href="#other_table_names-91"><span class="linenos">91</span></a> <span class="n">on</span> <span class="o">=</span> <span class="n">join</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;on&quot;</span><span class="p">)</span>
</span><span id="other_table_names-92"><a href="#other_table_names-92"><span class="linenos">92</span></a> <span class="k">return</span> <span class="n">exp</span><span class="o">.</span><span class="n">column_table_names</span><span class="p">(</span><span class="n">on</span><span class="p">,</span> <span class="n">join</span><span class="o">.</span><span class="n">alias_or_name</span><span class="p">)</span> <span class="k">if</span> <span class="n">on</span> <span class="k">else</span> <span class="nb">set</span><span class="p">()</span>
</span></pre></div>

File diff suppressed because one or more lines are too long

View file

@ -73,127 +73,133 @@
</span><span id="L-9"><a href="#L-9"><span class="linenos"> 9</span></a><span class="n">SELECT_ALL</span> <span class="o">=</span> <span class="nb">object</span><span class="p">()</span>
</span><span id="L-10"><a href="#L-10"><span class="linenos"> 10</span></a>
</span><span id="L-11"><a href="#L-11"><span class="linenos"> 11</span></a><span class="c1"># Selection to use if selection list is empty</span>
</span><span id="L-12"><a href="#L-12"><span class="linenos"> 12</span></a><span class="n">DEFAULT_SELECTION</span> <span class="o">=</span> <span class="k">lambda</span><span class="p">:</span> <span class="n">alias</span><span class="p">(</span><span class="s2">&quot;1&quot;</span><span class="p">,</span> <span class="s2">&quot;_&quot;</span><span class="p">)</span>
</span><span id="L-13"><a href="#L-13"><span class="linenos"> 13</span></a>
</span><span id="L-14"><a href="#L-14"><span class="linenos"> 14</span></a>
</span><span id="L-15"><a href="#L-15"><span class="linenos"> 15</span></a><span class="k">def</span> <span class="nf">pushdown_projections</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">schema</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">remove_unused_selections</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
</span><span id="L-16"><a href="#L-16"><span class="linenos"> 16</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-17"><a href="#L-17"><span class="linenos"> 17</span></a><span class="sd"> Rewrite sqlglot AST to remove unused columns projections.</span>
</span><span id="L-18"><a href="#L-18"><span class="linenos"> 18</span></a>
</span><span id="L-19"><a href="#L-19"><span class="linenos"> 19</span></a><span class="sd"> Example:</span>
</span><span id="L-20"><a href="#L-20"><span class="linenos"> 20</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="L-21"><a href="#L-21"><span class="linenos"> 21</span></a><span class="sd"> &gt;&gt;&gt; sql = &quot;SELECT y.a AS a FROM (SELECT x.a AS a, x.b AS b FROM x) AS y&quot;</span>
</span><span id="L-22"><a href="#L-22"><span class="linenos"> 22</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(sql)</span>
</span><span id="L-23"><a href="#L-23"><span class="linenos"> 23</span></a><span class="sd"> &gt;&gt;&gt; pushdown_projections(expression).sql()</span>
</span><span id="L-24"><a href="#L-24"><span class="linenos"> 24</span></a><span class="sd"> &#39;SELECT y.a AS a FROM (SELECT x.a AS a FROM x) AS y&#39;</span>
</span><span id="L-25"><a href="#L-25"><span class="linenos"> 25</span></a>
</span><span id="L-26"><a href="#L-26"><span class="linenos"> 26</span></a><span class="sd"> Args:</span>
</span><span id="L-27"><a href="#L-27"><span class="linenos"> 27</span></a><span class="sd"> expression (sqlglot.Expression): expression to optimize</span>
</span><span id="L-28"><a href="#L-28"><span class="linenos"> 28</span></a><span class="sd"> remove_unused_selections (bool): remove selects that are unused</span>
</span><span id="L-29"><a href="#L-29"><span class="linenos"> 29</span></a><span class="sd"> Returns:</span>
</span><span id="L-30"><a href="#L-30"><span class="linenos"> 30</span></a><span class="sd"> sqlglot.Expression: optimized expression</span>
</span><span id="L-31"><a href="#L-31"><span class="linenos"> 31</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-32"><a href="#L-32"><span class="linenos"> 32</span></a> <span class="c1"># Map of Scope to all columns being selected by outer queries.</span>
</span><span id="L-33"><a href="#L-33"><span class="linenos"> 33</span></a> <span class="n">schema</span> <span class="o">=</span> <span class="n">ensure_schema</span><span class="p">(</span><span class="n">schema</span><span class="p">)</span>
</span><span id="L-34"><a href="#L-34"><span class="linenos"> 34</span></a> <span class="n">source_column_alias_count</span> <span class="o">=</span> <span class="p">{}</span>
</span><span id="L-35"><a href="#L-35"><span class="linenos"> 35</span></a> <span class="n">referenced_columns</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="L-36"><a href="#L-36"><span class="linenos"> 36</span></a>
</span><span id="L-37"><a href="#L-37"><span class="linenos"> 37</span></a> <span class="c1"># We build the scope tree (which is traversed in DFS postorder), then iterate</span>
</span><span id="L-38"><a href="#L-38"><span class="linenos"> 38</span></a> <span class="c1"># over the result in reverse order. This should ensure that the set of selected</span>
</span><span id="L-39"><a href="#L-39"><span class="linenos"> 39</span></a> <span class="c1"># columns for a particular scope are completely build by the time we get to it.</span>
</span><span id="L-40"><a href="#L-40"><span class="linenos"> 40</span></a> <span class="k">for</span> <span class="n">scope</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">traverse_scope</span><span class="p">(</span><span class="n">expression</span><span class="p">)):</span>
</span><span id="L-41"><a href="#L-41"><span class="linenos"> 41</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="n">referenced_columns</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">})</span>
</span><span id="L-42"><a href="#L-42"><span class="linenos"> 42</span></a> <span class="n">alias_count</span> <span class="o">=</span> <span class="n">source_column_alias_count</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
</span><span id="L-43"><a href="#L-43"><span class="linenos"> 43</span></a>
</span><span id="L-44"><a href="#L-44"><span class="linenos"> 44</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;distinct&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">parent</span> <span class="ow">and</span> <span class="n">scope</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">pivots</span><span class="p">):</span>
</span><span id="L-45"><a href="#L-45"><span class="linenos"> 45</span></a> <span class="c1"># We can&#39;t remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if</span>
</span><span id="L-46"><a href="#L-46"><span class="linenos"> 46</span></a> <span class="c1"># we select from a pivoted source in the parent scope.</span>
</span><span id="L-47"><a href="#L-47"><span class="linenos"> 47</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">}</span>
</span><span id="L-48"><a href="#L-48"><span class="linenos"> 48</span></a>
</span><span id="L-49"><a href="#L-49"><span class="linenos"> 49</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Union</span><span class="p">):</span>
</span><span id="L-50"><a href="#L-50"><span class="linenos"> 50</span></a> <span class="n">left</span><span class="p">,</span> <span class="n">right</span> <span class="o">=</span> <span class="n">scope</span><span class="o">.</span><span class="n">union_scopes</span>
</span><span id="L-51"><a href="#L-51"><span class="linenos"> 51</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">left</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos"> 52</span></a>
</span><span id="L-53"><a href="#L-53"><span class="linenos"> 53</span></a> <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="L-54"><a href="#L-54"><span class="linenos"> 54</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="L-55"><a href="#L-55"><span class="linenos"> 55</span></a> <span class="k">elif</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="L-56"><a href="#L-56"><span class="linenos"> 56</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span>
</span><span id="L-57"><a href="#L-57"><span class="linenos"> 57</span></a> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">alias_or_name</span>
</span><span id="L-58"><a href="#L-58"><span class="linenos"> 58</span></a> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">select</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">)</span>
</span><span id="L-59"><a href="#L-59"><span class="linenos"> 59</span></a> <span class="k">if</span> <span class="n">SELECT_ALL</span> <span class="ow">in</span> <span class="n">parent_selections</span> <span class="ow">or</span> <span class="n">select</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">in</span> <span class="n">parent_selections</span>
</span><span id="L-60"><a href="#L-60"><span class="linenos"> 60</span></a> <span class="p">]</span>
</span><span id="L-61"><a href="#L-61"><span class="linenos"> 61</span></a>
</span><span id="L-62"><a href="#L-62"><span class="linenos"> 62</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Select</span><span class="p">):</span>
</span><span id="L-63"><a href="#L-63"><span class="linenos"> 63</span></a> <span class="k">if</span> <span class="n">remove_unused_selections</span><span class="p">:</span>
</span><span id="L-64"><a href="#L-64"><span class="linenos"> 64</span></a> <span class="n">_remove_unused_selections</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">parent_selections</span><span class="p">,</span> <span class="n">schema</span><span class="p">,</span> <span class="n">alias_count</span><span class="p">)</span>
</span><span id="L-65"><a href="#L-65"><span class="linenos"> 65</span></a>
</span><span id="L-66"><a href="#L-66"><span class="linenos"> 66</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">is_star</span><span class="p">:</span>
</span><span id="L-67"><a href="#L-67"><span class="linenos"> 67</span></a> <span class="k">continue</span>
</span><span id="L-68"><a href="#L-68"><span class="linenos"> 68</span></a>
</span><span id="L-69"><a href="#L-69"><span class="linenos"> 69</span></a> <span class="c1"># Group columns by source name</span>
</span><span id="L-70"><a href="#L-70"><span class="linenos"> 70</span></a> <span class="n">selects</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="L-71"><a href="#L-71"><span class="linenos"> 71</span></a> <span class="k">for</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">columns</span><span class="p">:</span>
</span><span id="L-72"><a href="#L-72"><span class="linenos"> 72</span></a> <span class="n">table_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">table</span>
</span><span id="L-73"><a href="#L-73"><span class="linenos"> 73</span></a> <span class="n">col_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">name</span>
</span><span id="L-74"><a href="#L-74"><span class="linenos"> 74</span></a> <span class="n">selects</span><span class="p">[</span><span class="n">table_name</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">col_name</span><span class="p">)</span>
</span><span id="L-75"><a href="#L-75"><span class="linenos"> 75</span></a>
</span><span id="L-76"><a href="#L-76"><span class="linenos"> 76</span></a> <span class="c1"># Push the selected columns down to the next scope</span>
</span><span id="L-77"><a href="#L-77"><span class="linenos"> 77</span></a> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">source</span><span class="p">)</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">selected_sources</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
</span><span id="L-78"><a href="#L-78"><span class="linenos"> 78</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">source</span><span class="p">,</span> <span class="n">Scope</span><span class="p">):</span>
</span><span id="L-79"><a href="#L-79"><span class="linenos"> 79</span></a> <span class="n">columns</span> <span class="o">=</span> <span class="n">selects</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="L-80"><a href="#L-80"><span class="linenos"> 80</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">source</span><span class="p">]</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">columns</span><span class="p">)</span>
</span><span id="L-81"><a href="#L-81"><span class="linenos"> 81</span></a>
</span><span id="L-82"><a href="#L-82"><span class="linenos"> 82</span></a> <span class="n">column_aliases</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">alias_column_names</span>
</span><span id="L-83"><a href="#L-83"><span class="linenos"> 83</span></a> <span class="k">if</span> <span class="n">column_aliases</span><span class="p">:</span>
</span><span id="L-84"><a href="#L-84"><span class="linenos"> 84</span></a> <span class="n">source_column_alias_count</span><span class="p">[</span><span class="n">source</span><span class="p">]</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">column_aliases</span><span class="p">)</span>
</span><span id="L-85"><a href="#L-85"><span class="linenos"> 85</span></a>
</span><span id="L-86"><a href="#L-86"><span class="linenos"> 86</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-12"><a href="#L-12"><span class="linenos"> 12</span></a><span class="n">DEFAULT_SELECTION</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">is_agg</span><span class="p">:</span> <span class="n">alias</span><span class="p">(</span>
</span><span id="L-13"><a href="#L-13"><span class="linenos"> 13</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">Max</span><span class="p">(</span><span class="n">this</span><span class="o">=</span><span class="n">exp</span><span class="o">.</span><span class="n">Literal</span><span class="o">.</span><span class="n">number</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span> <span class="k">if</span> <span class="n">is_agg</span> <span class="k">else</span> <span class="s2">&quot;1&quot;</span><span class="p">,</span> <span class="s2">&quot;_&quot;</span>
</span><span id="L-14"><a href="#L-14"><span class="linenos"> 14</span></a><span class="p">)</span>
</span><span id="L-15"><a href="#L-15"><span class="linenos"> 15</span></a>
</span><span id="L-16"><a href="#L-16"><span class="linenos"> 16</span></a>
</span><span id="L-17"><a href="#L-17"><span class="linenos"> 17</span></a><span class="k">def</span> <span class="nf">pushdown_projections</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">schema</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">remove_unused_selections</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
</span><span id="L-18"><a href="#L-18"><span class="linenos"> 18</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="L-19"><a href="#L-19"><span class="linenos"> 19</span></a><span class="sd"> Rewrite sqlglot AST to remove unused columns projections.</span>
</span><span id="L-20"><a href="#L-20"><span class="linenos"> 20</span></a>
</span><span id="L-21"><a href="#L-21"><span class="linenos"> 21</span></a><span class="sd"> Example:</span>
</span><span id="L-22"><a href="#L-22"><span class="linenos"> 22</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="L-23"><a href="#L-23"><span class="linenos"> 23</span></a><span class="sd"> &gt;&gt;&gt; sql = &quot;SELECT y.a AS a FROM (SELECT x.a AS a, x.b AS b FROM x) AS y&quot;</span>
</span><span id="L-24"><a href="#L-24"><span class="linenos"> 24</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(sql)</span>
</span><span id="L-25"><a href="#L-25"><span class="linenos"> 25</span></a><span class="sd"> &gt;&gt;&gt; pushdown_projections(expression).sql()</span>
</span><span id="L-26"><a href="#L-26"><span class="linenos"> 26</span></a><span class="sd"> &#39;SELECT y.a AS a FROM (SELECT x.a AS a FROM x) AS y&#39;</span>
</span><span id="L-27"><a href="#L-27"><span class="linenos"> 27</span></a>
</span><span id="L-28"><a href="#L-28"><span class="linenos"> 28</span></a><span class="sd"> Args:</span>
</span><span id="L-29"><a href="#L-29"><span class="linenos"> 29</span></a><span class="sd"> expression (sqlglot.Expression): expression to optimize</span>
</span><span id="L-30"><a href="#L-30"><span class="linenos"> 30</span></a><span class="sd"> remove_unused_selections (bool): remove selects that are unused</span>
</span><span id="L-31"><a href="#L-31"><span class="linenos"> 31</span></a><span class="sd"> Returns:</span>
</span><span id="L-32"><a href="#L-32"><span class="linenos"> 32</span></a><span class="sd"> sqlglot.Expression: optimized expression</span>
</span><span id="L-33"><a href="#L-33"><span class="linenos"> 33</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="L-34"><a href="#L-34"><span class="linenos"> 34</span></a> <span class="c1"># Map of Scope to all columns being selected by outer queries.</span>
</span><span id="L-35"><a href="#L-35"><span class="linenos"> 35</span></a> <span class="n">schema</span> <span class="o">=</span> <span class="n">ensure_schema</span><span class="p">(</span><span class="n">schema</span><span class="p">)</span>
</span><span id="L-36"><a href="#L-36"><span class="linenos"> 36</span></a> <span class="n">source_column_alias_count</span> <span class="o">=</span> <span class="p">{}</span>
</span><span id="L-37"><a href="#L-37"><span class="linenos"> 37</span></a> <span class="n">referenced_columns</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="L-38"><a href="#L-38"><span class="linenos"> 38</span></a>
</span><span id="L-39"><a href="#L-39"><span class="linenos"> 39</span></a> <span class="c1"># We build the scope tree (which is traversed in DFS postorder), then iterate</span>
</span><span id="L-40"><a href="#L-40"><span class="linenos"> 40</span></a> <span class="c1"># over the result in reverse order. This should ensure that the set of selected</span>
</span><span id="L-41"><a href="#L-41"><span class="linenos"> 41</span></a> <span class="c1"># columns for a particular scope are completely build by the time we get to it.</span>
</span><span id="L-42"><a href="#L-42"><span class="linenos"> 42</span></a> <span class="k">for</span> <span class="n">scope</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">traverse_scope</span><span class="p">(</span><span class="n">expression</span><span class="p">)):</span>
</span><span id="L-43"><a href="#L-43"><span class="linenos"> 43</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="n">referenced_columns</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">})</span>
</span><span id="L-44"><a href="#L-44"><span class="linenos"> 44</span></a> <span class="n">alias_count</span> <span class="o">=</span> <span class="n">source_column_alias_count</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
</span><span id="L-45"><a href="#L-45"><span class="linenos"> 45</span></a>
</span><span id="L-46"><a href="#L-46"><span class="linenos"> 46</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;distinct&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">parent</span> <span class="ow">and</span> <span class="n">scope</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">pivots</span><span class="p">):</span>
</span><span id="L-47"><a href="#L-47"><span class="linenos"> 47</span></a> <span class="c1"># We can&#39;t remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if</span>
</span><span id="L-48"><a href="#L-48"><span class="linenos"> 48</span></a> <span class="c1"># we select from a pivoted source in the parent scope.</span>
</span><span id="L-49"><a href="#L-49"><span class="linenos"> 49</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">}</span>
</span><span id="L-50"><a href="#L-50"><span class="linenos"> 50</span></a>
</span><span id="L-51"><a href="#L-51"><span class="linenos"> 51</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Union</span><span class="p">):</span>
</span><span id="L-52"><a href="#L-52"><span class="linenos"> 52</span></a> <span class="n">left</span><span class="p">,</span> <span class="n">right</span> <span class="o">=</span> <span class="n">scope</span><span class="o">.</span><span class="n">union_scopes</span>
</span><span id="L-53"><a href="#L-53"><span class="linenos"> 53</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">left</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="L-54"><a href="#L-54"><span class="linenos"> 54</span></a>
</span><span id="L-55"><a href="#L-55"><span class="linenos"> 55</span></a> <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="L-56"><a href="#L-56"><span class="linenos"> 56</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="L-57"><a href="#L-57"><span class="linenos"> 57</span></a> <span class="k">elif</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="L-58"><a href="#L-58"><span class="linenos"> 58</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span>
</span><span id="L-59"><a href="#L-59"><span class="linenos"> 59</span></a> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">alias_or_name</span>
</span><span id="L-60"><a href="#L-60"><span class="linenos"> 60</span></a> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">select</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">)</span>
</span><span id="L-61"><a href="#L-61"><span class="linenos"> 61</span></a> <span class="k">if</span> <span class="n">SELECT_ALL</span> <span class="ow">in</span> <span class="n">parent_selections</span> <span class="ow">or</span> <span class="n">select</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">in</span> <span class="n">parent_selections</span>
</span><span id="L-62"><a href="#L-62"><span class="linenos"> 62</span></a> <span class="p">]</span>
</span><span id="L-63"><a href="#L-63"><span class="linenos"> 63</span></a>
</span><span id="L-64"><a href="#L-64"><span class="linenos"> 64</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Select</span><span class="p">):</span>
</span><span id="L-65"><a href="#L-65"><span class="linenos"> 65</span></a> <span class="k">if</span> <span class="n">remove_unused_selections</span><span class="p">:</span>
</span><span id="L-66"><a href="#L-66"><span class="linenos"> 66</span></a> <span class="n">_remove_unused_selections</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">parent_selections</span><span class="p">,</span> <span class="n">schema</span><span class="p">,</span> <span class="n">alias_count</span><span class="p">)</span>
</span><span id="L-67"><a href="#L-67"><span class="linenos"> 67</span></a>
</span><span id="L-68"><a href="#L-68"><span class="linenos"> 68</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">is_star</span><span class="p">:</span>
</span><span id="L-69"><a href="#L-69"><span class="linenos"> 69</span></a> <span class="k">continue</span>
</span><span id="L-70"><a href="#L-70"><span class="linenos"> 70</span></a>
</span><span id="L-71"><a href="#L-71"><span class="linenos"> 71</span></a> <span class="c1"># Group columns by source name</span>
</span><span id="L-72"><a href="#L-72"><span class="linenos"> 72</span></a> <span class="n">selects</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="L-73"><a href="#L-73"><span class="linenos"> 73</span></a> <span class="k">for</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">columns</span><span class="p">:</span>
</span><span id="L-74"><a href="#L-74"><span class="linenos"> 74</span></a> <span class="n">table_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">table</span>
</span><span id="L-75"><a href="#L-75"><span class="linenos"> 75</span></a> <span class="n">col_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">name</span>
</span><span id="L-76"><a href="#L-76"><span class="linenos"> 76</span></a> <span class="n">selects</span><span class="p">[</span><span class="n">table_name</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">col_name</span><span class="p">)</span>
</span><span id="L-77"><a href="#L-77"><span class="linenos"> 77</span></a>
</span><span id="L-78"><a href="#L-78"><span class="linenos"> 78</span></a> <span class="c1"># Push the selected columns down to the next scope</span>
</span><span id="L-79"><a href="#L-79"><span class="linenos"> 79</span></a> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">source</span><span class="p">)</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">selected_sources</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
</span><span id="L-80"><a href="#L-80"><span class="linenos"> 80</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">source</span><span class="p">,</span> <span class="n">Scope</span><span class="p">):</span>
</span><span id="L-81"><a href="#L-81"><span class="linenos"> 81</span></a> <span class="n">columns</span> <span class="o">=</span> <span class="n">selects</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="L-82"><a href="#L-82"><span class="linenos"> 82</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">source</span><span class="p">]</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">columns</span><span class="p">)</span>
</span><span id="L-83"><a href="#L-83"><span class="linenos"> 83</span></a>
</span><span id="L-84"><a href="#L-84"><span class="linenos"> 84</span></a> <span class="n">column_aliases</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">alias_column_names</span>
</span><span id="L-85"><a href="#L-85"><span class="linenos"> 85</span></a> <span class="k">if</span> <span class="n">column_aliases</span><span class="p">:</span>
</span><span id="L-86"><a href="#L-86"><span class="linenos"> 86</span></a> <span class="n">source_column_alias_count</span><span class="p">[</span><span class="n">source</span><span class="p">]</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">column_aliases</span><span class="p">)</span>
</span><span id="L-87"><a href="#L-87"><span class="linenos"> 87</span></a>
</span><span id="L-88"><a href="#L-88"><span class="linenos"> 88</span></a>
</span><span id="L-89"><a href="#L-89"><span class="linenos"> 89</span></a><span class="k">def</span> <span class="nf">_remove_unused_selections</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">parent_selections</span><span class="p">,</span> <span class="n">schema</span><span class="p">,</span> <span class="n">alias_count</span><span class="p">):</span>
</span><span id="L-90"><a href="#L-90"><span class="linenos"> 90</span></a> <span class="n">order</span> <span class="o">=</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;order&quot;</span><span class="p">)</span>
</span><span id="L-91"><a href="#L-91"><span class="linenos"> 91</span></a>
</span><span id="L-92"><a href="#L-92"><span class="linenos"> 92</span></a> <span class="k">if</span> <span class="n">order</span><span class="p">:</span>
</span><span id="L-93"><a href="#L-93"><span class="linenos"> 93</span></a> <span class="c1"># Assume columns without a qualified table are references to output columns</span>
</span><span id="L-94"><a href="#L-94"><span class="linenos"> 94</span></a> <span class="n">order_refs</span> <span class="o">=</span> <span class="p">{</span><span class="n">c</span><span class="o">.</span><span class="n">name</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="n">order</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Column</span><span class="p">)</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">c</span><span class="o">.</span><span class="n">table</span><span class="p">}</span>
</span><span id="L-95"><a href="#L-95"><span class="linenos"> 95</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-96"><a href="#L-96"><span class="linenos"> 96</span></a> <span class="n">order_refs</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="L-97"><a href="#L-97"><span class="linenos"> 97</span></a>
</span><span id="L-98"><a href="#L-98"><span class="linenos"> 98</span></a> <span class="n">new_selections</span> <span class="o">=</span> <span class="p">[]</span>
</span><span id="L-99"><a href="#L-99"><span class="linenos"> 99</span></a> <span class="n">removed</span> <span class="o">=</span> <span class="kc">False</span>
</span><span id="L-100"><a href="#L-100"><span class="linenos">100</span></a> <span class="n">star</span> <span class="o">=</span> <span class="kc">False</span>
</span><span id="L-101"><a href="#L-101"><span class="linenos">101</span></a>
</span><span id="L-102"><a href="#L-102"><span class="linenos">102</span></a> <span class="n">select_all</span> <span class="o">=</span> <span class="n">SELECT_ALL</span> <span class="ow">in</span> <span class="n">parent_selections</span>
</span><span id="L-103"><a href="#L-103"><span class="linenos">103</span></a>
</span><span id="L-104"><a href="#L-104"><span class="linenos">104</span></a> <span class="k">for</span> <span class="n">selection</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">:</span>
</span><span id="L-105"><a href="#L-105"><span class="linenos">105</span></a> <span class="n">name</span> <span class="o">=</span> <span class="n">selection</span><span class="o">.</span><span class="n">alias_or_name</span>
</span><span id="L-88"><a href="#L-88"><span class="linenos"> 88</span></a> <span class="k">return</span> <span class="n">expression</span>
</span><span id="L-89"><a href="#L-89"><span class="linenos"> 89</span></a>
</span><span id="L-90"><a href="#L-90"><span class="linenos"> 90</span></a>
</span><span id="L-91"><a href="#L-91"><span class="linenos"> 91</span></a><span class="k">def</span> <span class="nf">_remove_unused_selections</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">parent_selections</span><span class="p">,</span> <span class="n">schema</span><span class="p">,</span> <span class="n">alias_count</span><span class="p">):</span>
</span><span id="L-92"><a href="#L-92"><span class="linenos"> 92</span></a> <span class="n">order</span> <span class="o">=</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;order&quot;</span><span class="p">)</span>
</span><span id="L-93"><a href="#L-93"><span class="linenos"> 93</span></a>
</span><span id="L-94"><a href="#L-94"><span class="linenos"> 94</span></a> <span class="k">if</span> <span class="n">order</span><span class="p">:</span>
</span><span id="L-95"><a href="#L-95"><span class="linenos"> 95</span></a> <span class="c1"># Assume columns without a qualified table are references to output columns</span>
</span><span id="L-96"><a href="#L-96"><span class="linenos"> 96</span></a> <span class="n">order_refs</span> <span class="o">=</span> <span class="p">{</span><span class="n">c</span><span class="o">.</span><span class="n">name</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="n">order</span><span class="o">.</span><span class="n">find_all</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Column</span><span class="p">)</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">c</span><span class="o">.</span><span class="n">table</span><span class="p">}</span>
</span><span id="L-97"><a href="#L-97"><span class="linenos"> 97</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-98"><a href="#L-98"><span class="linenos"> 98</span></a> <span class="n">order_refs</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="L-99"><a href="#L-99"><span class="linenos"> 99</span></a>
</span><span id="L-100"><a href="#L-100"><span class="linenos">100</span></a> <span class="n">new_selections</span> <span class="o">=</span> <span class="p">[]</span>
</span><span id="L-101"><a href="#L-101"><span class="linenos">101</span></a> <span class="n">removed</span> <span class="o">=</span> <span class="kc">False</span>
</span><span id="L-102"><a href="#L-102"><span class="linenos">102</span></a> <span class="n">star</span> <span class="o">=</span> <span class="kc">False</span>
</span><span id="L-103"><a href="#L-103"><span class="linenos">103</span></a> <span class="n">is_agg</span> <span class="o">=</span> <span class="kc">False</span>
</span><span id="L-104"><a href="#L-104"><span class="linenos">104</span></a>
</span><span id="L-105"><a href="#L-105"><span class="linenos">105</span></a> <span class="n">select_all</span> <span class="o">=</span> <span class="n">SELECT_ALL</span> <span class="ow">in</span> <span class="n">parent_selections</span>
</span><span id="L-106"><a href="#L-106"><span class="linenos">106</span></a>
</span><span id="L-107"><a href="#L-107"><span class="linenos">107</span></a> <span class="k">if</span> <span class="n">select_all</span> <span class="ow">or</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">parent_selections</span> <span class="ow">or</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">order_refs</span> <span class="ow">or</span> <span class="n">alias_count</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
</span><span id="L-108"><a href="#L-108"><span class="linenos">108</span></a> <span class="n">new_selections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">selection</span><span class="p">)</span>
</span><span id="L-109"><a href="#L-109"><span class="linenos">109</span></a> <span class="n">alias_count</span> <span class="o">-=</span> <span class="mi">1</span>
</span><span id="L-110"><a href="#L-110"><span class="linenos">110</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-111"><a href="#L-111"><span class="linenos">111</span></a> <span class="k">if</span> <span class="n">selection</span><span class="o">.</span><span class="n">is_star</span><span class="p">:</span>
</span><span id="L-112"><a href="#L-112"><span class="linenos">112</span></a> <span class="n">star</span> <span class="o">=</span> <span class="kc">True</span>
</span><span id="L-113"><a href="#L-113"><span class="linenos">113</span></a> <span class="n">removed</span> <span class="o">=</span> <span class="kc">True</span>
</span><span id="L-114"><a href="#L-114"><span class="linenos">114</span></a>
</span><span id="L-115"><a href="#L-115"><span class="linenos">115</span></a> <span class="k">if</span> <span class="n">star</span><span class="p">:</span>
</span><span id="L-116"><a href="#L-116"><span class="linenos">116</span></a> <span class="n">resolver</span> <span class="o">=</span> <span class="n">Resolver</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">schema</span><span class="p">)</span>
</span><span id="L-117"><a href="#L-117"><span class="linenos">117</span></a> <span class="n">names</span> <span class="o">=</span> <span class="p">{</span><span class="n">s</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">new_selections</span><span class="p">}</span>
</span><span id="L-118"><a href="#L-118"><span class="linenos">118</span></a>
</span><span id="L-119"><a href="#L-119"><span class="linenos">119</span></a> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">parent_selections</span><span class="p">):</span>
</span><span id="L-120"><a href="#L-120"><span class="linenos">120</span></a> <span class="k">if</span> <span class="n">name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">names</span><span class="p">:</span>
</span><span id="L-121"><a href="#L-121"><span class="linenos">121</span></a> <span class="n">new_selections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
</span><span id="L-122"><a href="#L-122"><span class="linenos">122</span></a> <span class="n">alias</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">column</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">table</span><span class="o">=</span><span class="n">resolver</span><span class="o">.</span><span class="n">get_table</span><span class="p">(</span><span class="n">name</span><span class="p">)),</span> <span class="n">name</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-123"><a href="#L-123"><span class="linenos">123</span></a> <span class="p">)</span>
</span><span id="L-107"><a href="#L-107"><span class="linenos">107</span></a> <span class="k">for</span> <span class="n">selection</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">:</span>
</span><span id="L-108"><a href="#L-108"><span class="linenos">108</span></a> <span class="n">name</span> <span class="o">=</span> <span class="n">selection</span><span class="o">.</span><span class="n">alias_or_name</span>
</span><span id="L-109"><a href="#L-109"><span class="linenos">109</span></a>
</span><span id="L-110"><a href="#L-110"><span class="linenos">110</span></a> <span class="k">if</span> <span class="n">select_all</span> <span class="ow">or</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">parent_selections</span> <span class="ow">or</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">order_refs</span> <span class="ow">or</span> <span class="n">alias_count</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
</span><span id="L-111"><a href="#L-111"><span class="linenos">111</span></a> <span class="n">new_selections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">selection</span><span class="p">)</span>
</span><span id="L-112"><a href="#L-112"><span class="linenos">112</span></a> <span class="n">alias_count</span> <span class="o">-=</span> <span class="mi">1</span>
</span><span id="L-113"><a href="#L-113"><span class="linenos">113</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-114"><a href="#L-114"><span class="linenos">114</span></a> <span class="k">if</span> <span class="n">selection</span><span class="o">.</span><span class="n">is_star</span><span class="p">:</span>
</span><span id="L-115"><a href="#L-115"><span class="linenos">115</span></a> <span class="n">star</span> <span class="o">=</span> <span class="kc">True</span>
</span><span id="L-116"><a href="#L-116"><span class="linenos">116</span></a> <span class="n">removed</span> <span class="o">=</span> <span class="kc">True</span>
</span><span id="L-117"><a href="#L-117"><span class="linenos">117</span></a>
</span><span id="L-118"><a href="#L-118"><span class="linenos">118</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">is_agg</span> <span class="ow">and</span> <span class="n">selection</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">AggFunc</span><span class="p">):</span>
</span><span id="L-119"><a href="#L-119"><span class="linenos">119</span></a> <span class="n">is_agg</span> <span class="o">=</span> <span class="kc">True</span>
</span><span id="L-120"><a href="#L-120"><span class="linenos">120</span></a>
</span><span id="L-121"><a href="#L-121"><span class="linenos">121</span></a> <span class="k">if</span> <span class="n">star</span><span class="p">:</span>
</span><span id="L-122"><a href="#L-122"><span class="linenos">122</span></a> <span class="n">resolver</span> <span class="o">=</span> <span class="n">Resolver</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">schema</span><span class="p">)</span>
</span><span id="L-123"><a href="#L-123"><span class="linenos">123</span></a> <span class="n">names</span> <span class="o">=</span> <span class="p">{</span><span class="n">s</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">new_selections</span><span class="p">}</span>
</span><span id="L-124"><a href="#L-124"><span class="linenos">124</span></a>
</span><span id="L-125"><a href="#L-125"><span class="linenos">125</span></a> <span class="c1"># If there are no remaining selections, just select a single constant</span>
</span><span id="L-126"><a href="#L-126"><span class="linenos">126</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">new_selections</span><span class="p">:</span>
</span><span id="L-127"><a href="#L-127"><span class="linenos">127</span></a> <span class="n">new_selections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">DEFAULT_SELECTION</span><span class="p">())</span>
</span><span id="L-128"><a href="#L-128"><span class="linenos">128</span></a>
</span><span id="L-129"><a href="#L-129"><span class="linenos">129</span></a> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="o">*</span><span class="n">new_selections</span><span class="p">,</span> <span class="n">append</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-125"><a href="#L-125"><span class="linenos">125</span></a> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">parent_selections</span><span class="p">):</span>
</span><span id="L-126"><a href="#L-126"><span class="linenos">126</span></a> <span class="k">if</span> <span class="n">name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">names</span><span class="p">:</span>
</span><span id="L-127"><a href="#L-127"><span class="linenos">127</span></a> <span class="n">new_selections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
</span><span id="L-128"><a href="#L-128"><span class="linenos">128</span></a> <span class="n">alias</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">column</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">table</span><span class="o">=</span><span class="n">resolver</span><span class="o">.</span><span class="n">get_table</span><span class="p">(</span><span class="n">name</span><span class="p">)),</span> <span class="n">name</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-129"><a href="#L-129"><span class="linenos">129</span></a> <span class="p">)</span>
</span><span id="L-130"><a href="#L-130"><span class="linenos">130</span></a>
</span><span id="L-131"><a href="#L-131"><span class="linenos">131</span></a> <span class="k">if</span> <span class="n">removed</span><span class="p">:</span>
</span><span id="L-132"><a href="#L-132"><span class="linenos">132</span></a> <span class="n">scope</span><span class="o">.</span><span class="n">clear_cache</span><span class="p">()</span>
</span><span id="L-131"><a href="#L-131"><span class="linenos">131</span></a> <span class="c1"># If there are no remaining selections, just select a single constant</span>
</span><span id="L-132"><a href="#L-132"><span class="linenos">132</span></a> <span class="k">if</span> <span class="ow">not</span> <span class="n">new_selections</span><span class="p">:</span>
</span><span id="L-133"><a href="#L-133"><span class="linenos">133</span></a> <span class="n">new_selections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">DEFAULT_SELECTION</span><span class="p">(</span><span class="n">is_agg</span><span class="p">))</span>
</span><span id="L-134"><a href="#L-134"><span class="linenos">134</span></a>
</span><span id="L-135"><a href="#L-135"><span class="linenos">135</span></a> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="o">*</span><span class="n">new_selections</span><span class="p">,</span> <span class="n">append</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</span><span id="L-136"><a href="#L-136"><span class="linenos">136</span></a>
</span><span id="L-137"><a href="#L-137"><span class="linenos">137</span></a> <span class="k">if</span> <span class="n">removed</span><span class="p">:</span>
</span><span id="L-138"><a href="#L-138"><span class="linenos">138</span></a> <span class="n">scope</span><span class="o">.</span><span class="n">clear_cache</span><span class="p">()</span>
</span></pre></div>
@ -215,13 +221,15 @@
<div class="attr function">
<span class="def">def</span>
<span class="name">DEFAULT_SELECTION</span><span class="signature pdoc-code condensed">(<span class="return-annotation">):</span></span>
<span class="name">DEFAULT_SELECTION</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">is_agg</span></span><span class="return-annotation">):</span></span>
<label class="view-source-button" for="DEFAULT_SELECTION-view-source"><span>View Source</span></label>
</div>
<a class="headerlink" href="#DEFAULT_SELECTION"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="DEFAULT_SELECTION-13"><a href="#DEFAULT_SELECTION-13"><span class="linenos">13</span></a><span class="n">DEFAULT_SELECTION</span> <span class="o">=</span> <span class="k">lambda</span><span class="p">:</span> <span class="n">alias</span><span class="p">(</span><span class="s2">&quot;1&quot;</span><span class="p">,</span> <span class="s2">&quot;_&quot;</span><span class="p">)</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="DEFAULT_SELECTION-13"><a href="#DEFAULT_SELECTION-13"><span class="linenos">13</span></a><span class="n">DEFAULT_SELECTION</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">is_agg</span><span class="p">:</span> <span class="n">alias</span><span class="p">(</span>
</span><span id="DEFAULT_SELECTION-14"><a href="#DEFAULT_SELECTION-14"><span class="linenos">14</span></a> <span class="n">exp</span><span class="o">.</span><span class="n">Max</span><span class="p">(</span><span class="n">this</span><span class="o">=</span><span class="n">exp</span><span class="o">.</span><span class="n">Literal</span><span class="o">.</span><span class="n">number</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span> <span class="k">if</span> <span class="n">is_agg</span> <span class="k">else</span> <span class="s2">&quot;1&quot;</span><span class="p">,</span> <span class="s2">&quot;_&quot;</span>
</span><span id="DEFAULT_SELECTION-15"><a href="#DEFAULT_SELECTION-15"><span class="linenos">15</span></a><span class="p">)</span>
</span></pre></div>
@ -239,78 +247,78 @@
</div>
<a class="headerlink" href="#pushdown_projections"></a>
<div class="pdoc-code codehilite"><pre><span></span><span id="pushdown_projections-16"><a href="#pushdown_projections-16"><span class="linenos">16</span></a><span class="k">def</span> <span class="nf">pushdown_projections</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">schema</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">remove_unused_selections</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
</span><span id="pushdown_projections-17"><a href="#pushdown_projections-17"><span class="linenos">17</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="pushdown_projections-18"><a href="#pushdown_projections-18"><span class="linenos">18</span></a><span class="sd"> Rewrite sqlglot AST to remove unused columns projections.</span>
</span><span id="pushdown_projections-19"><a href="#pushdown_projections-19"><span class="linenos">19</span></a>
</span><span id="pushdown_projections-20"><a href="#pushdown_projections-20"><span class="linenos">20</span></a><span class="sd"> Example:</span>
</span><span id="pushdown_projections-21"><a href="#pushdown_projections-21"><span class="linenos">21</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="pushdown_projections-22"><a href="#pushdown_projections-22"><span class="linenos">22</span></a><span class="sd"> &gt;&gt;&gt; sql = &quot;SELECT y.a AS a FROM (SELECT x.a AS a, x.b AS b FROM x) AS y&quot;</span>
</span><span id="pushdown_projections-23"><a href="#pushdown_projections-23"><span class="linenos">23</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(sql)</span>
</span><span id="pushdown_projections-24"><a href="#pushdown_projections-24"><span class="linenos">24</span></a><span class="sd"> &gt;&gt;&gt; pushdown_projections(expression).sql()</span>
</span><span id="pushdown_projections-25"><a href="#pushdown_projections-25"><span class="linenos">25</span></a><span class="sd"> &#39;SELECT y.a AS a FROM (SELECT x.a AS a FROM x) AS y&#39;</span>
</span><span id="pushdown_projections-26"><a href="#pushdown_projections-26"><span class="linenos">26</span></a>
</span><span id="pushdown_projections-27"><a href="#pushdown_projections-27"><span class="linenos">27</span></a><span class="sd"> Args:</span>
</span><span id="pushdown_projections-28"><a href="#pushdown_projections-28"><span class="linenos">28</span></a><span class="sd"> expression (sqlglot.Expression): expression to optimize</span>
</span><span id="pushdown_projections-29"><a href="#pushdown_projections-29"><span class="linenos">29</span></a><span class="sd"> remove_unused_selections (bool): remove selects that are unused</span>
</span><span id="pushdown_projections-30"><a href="#pushdown_projections-30"><span class="linenos">30</span></a><span class="sd"> Returns:</span>
</span><span id="pushdown_projections-31"><a href="#pushdown_projections-31"><span class="linenos">31</span></a><span class="sd"> sqlglot.Expression: optimized expression</span>
</span><span id="pushdown_projections-32"><a href="#pushdown_projections-32"><span class="linenos">32</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="pushdown_projections-33"><a href="#pushdown_projections-33"><span class="linenos">33</span></a> <span class="c1"># Map of Scope to all columns being selected by outer queries.</span>
</span><span id="pushdown_projections-34"><a href="#pushdown_projections-34"><span class="linenos">34</span></a> <span class="n">schema</span> <span class="o">=</span> <span class="n">ensure_schema</span><span class="p">(</span><span class="n">schema</span><span class="p">)</span>
</span><span id="pushdown_projections-35"><a href="#pushdown_projections-35"><span class="linenos">35</span></a> <span class="n">source_column_alias_count</span> <span class="o">=</span> <span class="p">{}</span>
</span><span id="pushdown_projections-36"><a href="#pushdown_projections-36"><span class="linenos">36</span></a> <span class="n">referenced_columns</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="pushdown_projections-37"><a href="#pushdown_projections-37"><span class="linenos">37</span></a>
</span><span id="pushdown_projections-38"><a href="#pushdown_projections-38"><span class="linenos">38</span></a> <span class="c1"># We build the scope tree (which is traversed in DFS postorder), then iterate</span>
</span><span id="pushdown_projections-39"><a href="#pushdown_projections-39"><span class="linenos">39</span></a> <span class="c1"># over the result in reverse order. This should ensure that the set of selected</span>
</span><span id="pushdown_projections-40"><a href="#pushdown_projections-40"><span class="linenos">40</span></a> <span class="c1"># columns for a particular scope are completely build by the time we get to it.</span>
</span><span id="pushdown_projections-41"><a href="#pushdown_projections-41"><span class="linenos">41</span></a> <span class="k">for</span> <span class="n">scope</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">traverse_scope</span><span class="p">(</span><span class="n">expression</span><span class="p">)):</span>
</span><span id="pushdown_projections-42"><a href="#pushdown_projections-42"><span class="linenos">42</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="n">referenced_columns</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">})</span>
</span><span id="pushdown_projections-43"><a href="#pushdown_projections-43"><span class="linenos">43</span></a> <span class="n">alias_count</span> <span class="o">=</span> <span class="n">source_column_alias_count</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
</span><span id="pushdown_projections-44"><a href="#pushdown_projections-44"><span class="linenos">44</span></a>
</span><span id="pushdown_projections-45"><a href="#pushdown_projections-45"><span class="linenos">45</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;distinct&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">parent</span> <span class="ow">and</span> <span class="n">scope</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">pivots</span><span class="p">):</span>
</span><span id="pushdown_projections-46"><a href="#pushdown_projections-46"><span class="linenos">46</span></a> <span class="c1"># We can&#39;t remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if</span>
</span><span id="pushdown_projections-47"><a href="#pushdown_projections-47"><span class="linenos">47</span></a> <span class="c1"># we select from a pivoted source in the parent scope.</span>
</span><span id="pushdown_projections-48"><a href="#pushdown_projections-48"><span class="linenos">48</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">}</span>
</span><span id="pushdown_projections-49"><a href="#pushdown_projections-49"><span class="linenos">49</span></a>
</span><span id="pushdown_projections-50"><a href="#pushdown_projections-50"><span class="linenos">50</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Union</span><span class="p">):</span>
</span><span id="pushdown_projections-51"><a href="#pushdown_projections-51"><span class="linenos">51</span></a> <span class="n">left</span><span class="p">,</span> <span class="n">right</span> <span class="o">=</span> <span class="n">scope</span><span class="o">.</span><span class="n">union_scopes</span>
</span><span id="pushdown_projections-52"><a href="#pushdown_projections-52"><span class="linenos">52</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">left</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="pushdown_projections-53"><a href="#pushdown_projections-53"><span class="linenos">53</span></a>
</span><span id="pushdown_projections-54"><a href="#pushdown_projections-54"><span class="linenos">54</span></a> <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="pushdown_projections-55"><a href="#pushdown_projections-55"><span class="linenos">55</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="pushdown_projections-56"><a href="#pushdown_projections-56"><span class="linenos">56</span></a> <span class="k">elif</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="pushdown_projections-57"><a href="#pushdown_projections-57"><span class="linenos">57</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span>
</span><span id="pushdown_projections-58"><a href="#pushdown_projections-58"><span class="linenos">58</span></a> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">alias_or_name</span>
</span><span id="pushdown_projections-59"><a href="#pushdown_projections-59"><span class="linenos">59</span></a> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">select</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">)</span>
</span><span id="pushdown_projections-60"><a href="#pushdown_projections-60"><span class="linenos">60</span></a> <span class="k">if</span> <span class="n">SELECT_ALL</span> <span class="ow">in</span> <span class="n">parent_selections</span> <span class="ow">or</span> <span class="n">select</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">in</span> <span class="n">parent_selections</span>
</span><span id="pushdown_projections-61"><a href="#pushdown_projections-61"><span class="linenos">61</span></a> <span class="p">]</span>
</span><span id="pushdown_projections-62"><a href="#pushdown_projections-62"><span class="linenos">62</span></a>
</span><span id="pushdown_projections-63"><a href="#pushdown_projections-63"><span class="linenos">63</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Select</span><span class="p">):</span>
</span><span id="pushdown_projections-64"><a href="#pushdown_projections-64"><span class="linenos">64</span></a> <span class="k">if</span> <span class="n">remove_unused_selections</span><span class="p">:</span>
</span><span id="pushdown_projections-65"><a href="#pushdown_projections-65"><span class="linenos">65</span></a> <span class="n">_remove_unused_selections</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">parent_selections</span><span class="p">,</span> <span class="n">schema</span><span class="p">,</span> <span class="n">alias_count</span><span class="p">)</span>
</span><span id="pushdown_projections-66"><a href="#pushdown_projections-66"><span class="linenos">66</span></a>
</span><span id="pushdown_projections-67"><a href="#pushdown_projections-67"><span class="linenos">67</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">is_star</span><span class="p">:</span>
</span><span id="pushdown_projections-68"><a href="#pushdown_projections-68"><span class="linenos">68</span></a> <span class="k">continue</span>
</span><span id="pushdown_projections-69"><a href="#pushdown_projections-69"><span class="linenos">69</span></a>
</span><span id="pushdown_projections-70"><a href="#pushdown_projections-70"><span class="linenos">70</span></a> <span class="c1"># Group columns by source name</span>
</span><span id="pushdown_projections-71"><a href="#pushdown_projections-71"><span class="linenos">71</span></a> <span class="n">selects</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="pushdown_projections-72"><a href="#pushdown_projections-72"><span class="linenos">72</span></a> <span class="k">for</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">columns</span><span class="p">:</span>
</span><span id="pushdown_projections-73"><a href="#pushdown_projections-73"><span class="linenos">73</span></a> <span class="n">table_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">table</span>
</span><span id="pushdown_projections-74"><a href="#pushdown_projections-74"><span class="linenos">74</span></a> <span class="n">col_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">name</span>
</span><span id="pushdown_projections-75"><a href="#pushdown_projections-75"><span class="linenos">75</span></a> <span class="n">selects</span><span class="p">[</span><span class="n">table_name</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">col_name</span><span class="p">)</span>
</span><span id="pushdown_projections-76"><a href="#pushdown_projections-76"><span class="linenos">76</span></a>
</span><span id="pushdown_projections-77"><a href="#pushdown_projections-77"><span class="linenos">77</span></a> <span class="c1"># Push the selected columns down to the next scope</span>
</span><span id="pushdown_projections-78"><a href="#pushdown_projections-78"><span class="linenos">78</span></a> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">source</span><span class="p">)</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">selected_sources</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
</span><span id="pushdown_projections-79"><a href="#pushdown_projections-79"><span class="linenos">79</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">source</span><span class="p">,</span> <span class="n">Scope</span><span class="p">):</span>
</span><span id="pushdown_projections-80"><a href="#pushdown_projections-80"><span class="linenos">80</span></a> <span class="n">columns</span> <span class="o">=</span> <span class="n">selects</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="pushdown_projections-81"><a href="#pushdown_projections-81"><span class="linenos">81</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">source</span><span class="p">]</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">columns</span><span class="p">)</span>
</span><span id="pushdown_projections-82"><a href="#pushdown_projections-82"><span class="linenos">82</span></a>
</span><span id="pushdown_projections-83"><a href="#pushdown_projections-83"><span class="linenos">83</span></a> <span class="n">column_aliases</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">alias_column_names</span>
</span><span id="pushdown_projections-84"><a href="#pushdown_projections-84"><span class="linenos">84</span></a> <span class="k">if</span> <span class="n">column_aliases</span><span class="p">:</span>
</span><span id="pushdown_projections-85"><a href="#pushdown_projections-85"><span class="linenos">85</span></a> <span class="n">source_column_alias_count</span><span class="p">[</span><span class="n">source</span><span class="p">]</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">column_aliases</span><span class="p">)</span>
</span><span id="pushdown_projections-86"><a href="#pushdown_projections-86"><span class="linenos">86</span></a>
</span><span id="pushdown_projections-87"><a href="#pushdown_projections-87"><span class="linenos">87</span></a> <span class="k">return</span> <span class="n">expression</span>
<div class="pdoc-code codehilite"><pre><span></span><span id="pushdown_projections-18"><a href="#pushdown_projections-18"><span class="linenos">18</span></a><span class="k">def</span> <span class="nf">pushdown_projections</span><span class="p">(</span><span class="n">expression</span><span class="p">,</span> <span class="n">schema</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">remove_unused_selections</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
</span><span id="pushdown_projections-19"><a href="#pushdown_projections-19"><span class="linenos">19</span></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="pushdown_projections-20"><a href="#pushdown_projections-20"><span class="linenos">20</span></a><span class="sd"> Rewrite sqlglot AST to remove unused columns projections.</span>
</span><span id="pushdown_projections-21"><a href="#pushdown_projections-21"><span class="linenos">21</span></a>
</span><span id="pushdown_projections-22"><a href="#pushdown_projections-22"><span class="linenos">22</span></a><span class="sd"> Example:</span>
</span><span id="pushdown_projections-23"><a href="#pushdown_projections-23"><span class="linenos">23</span></a><span class="sd"> &gt;&gt;&gt; import sqlglot</span>
</span><span id="pushdown_projections-24"><a href="#pushdown_projections-24"><span class="linenos">24</span></a><span class="sd"> &gt;&gt;&gt; sql = &quot;SELECT y.a AS a FROM (SELECT x.a AS a, x.b AS b FROM x) AS y&quot;</span>
</span><span id="pushdown_projections-25"><a href="#pushdown_projections-25"><span class="linenos">25</span></a><span class="sd"> &gt;&gt;&gt; expression = sqlglot.parse_one(sql)</span>
</span><span id="pushdown_projections-26"><a href="#pushdown_projections-26"><span class="linenos">26</span></a><span class="sd"> &gt;&gt;&gt; pushdown_projections(expression).sql()</span>
</span><span id="pushdown_projections-27"><a href="#pushdown_projections-27"><span class="linenos">27</span></a><span class="sd"> &#39;SELECT y.a AS a FROM (SELECT x.a AS a FROM x) AS y&#39;</span>
</span><span id="pushdown_projections-28"><a href="#pushdown_projections-28"><span class="linenos">28</span></a>
</span><span id="pushdown_projections-29"><a href="#pushdown_projections-29"><span class="linenos">29</span></a><span class="sd"> Args:</span>
</span><span id="pushdown_projections-30"><a href="#pushdown_projections-30"><span class="linenos">30</span></a><span class="sd"> expression (sqlglot.Expression): expression to optimize</span>
</span><span id="pushdown_projections-31"><a href="#pushdown_projections-31"><span class="linenos">31</span></a><span class="sd"> remove_unused_selections (bool): remove selects that are unused</span>
</span><span id="pushdown_projections-32"><a href="#pushdown_projections-32"><span class="linenos">32</span></a><span class="sd"> Returns:</span>
</span><span id="pushdown_projections-33"><a href="#pushdown_projections-33"><span class="linenos">33</span></a><span class="sd"> sqlglot.Expression: optimized expression</span>
</span><span id="pushdown_projections-34"><a href="#pushdown_projections-34"><span class="linenos">34</span></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="pushdown_projections-35"><a href="#pushdown_projections-35"><span class="linenos">35</span></a> <span class="c1"># Map of Scope to all columns being selected by outer queries.</span>
</span><span id="pushdown_projections-36"><a href="#pushdown_projections-36"><span class="linenos">36</span></a> <span class="n">schema</span> <span class="o">=</span> <span class="n">ensure_schema</span><span class="p">(</span><span class="n">schema</span><span class="p">)</span>
</span><span id="pushdown_projections-37"><a href="#pushdown_projections-37"><span class="linenos">37</span></a> <span class="n">source_column_alias_count</span> <span class="o">=</span> <span class="p">{}</span>
</span><span id="pushdown_projections-38"><a href="#pushdown_projections-38"><span class="linenos">38</span></a> <span class="n">referenced_columns</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="pushdown_projections-39"><a href="#pushdown_projections-39"><span class="linenos">39</span></a>
</span><span id="pushdown_projections-40"><a href="#pushdown_projections-40"><span class="linenos">40</span></a> <span class="c1"># We build the scope tree (which is traversed in DFS postorder), then iterate</span>
</span><span id="pushdown_projections-41"><a href="#pushdown_projections-41"><span class="linenos">41</span></a> <span class="c1"># over the result in reverse order. This should ensure that the set of selected</span>
</span><span id="pushdown_projections-42"><a href="#pushdown_projections-42"><span class="linenos">42</span></a> <span class="c1"># columns for a particular scope are completely build by the time we get to it.</span>
</span><span id="pushdown_projections-43"><a href="#pushdown_projections-43"><span class="linenos">43</span></a> <span class="k">for</span> <span class="n">scope</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">traverse_scope</span><span class="p">(</span><span class="n">expression</span><span class="p">)):</span>
</span><span id="pushdown_projections-44"><a href="#pushdown_projections-44"><span class="linenos">44</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="n">referenced_columns</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">})</span>
</span><span id="pushdown_projections-45"><a href="#pushdown_projections-45"><span class="linenos">45</span></a> <span class="n">alias_count</span> <span class="o">=</span> <span class="n">source_column_alias_count</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
</span><span id="pushdown_projections-46"><a href="#pushdown_projections-46"><span class="linenos">46</span></a>
</span><span id="pushdown_projections-47"><a href="#pushdown_projections-47"><span class="linenos">47</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">args</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;distinct&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">parent</span> <span class="ow">and</span> <span class="n">scope</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">pivots</span><span class="p">):</span>
</span><span id="pushdown_projections-48"><a href="#pushdown_projections-48"><span class="linenos">48</span></a> <span class="c1"># We can&#39;t remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if</span>
</span><span id="pushdown_projections-49"><a href="#pushdown_projections-49"><span class="linenos">49</span></a> <span class="c1"># we select from a pivoted source in the parent scope.</span>
</span><span id="pushdown_projections-50"><a href="#pushdown_projections-50"><span class="linenos">50</span></a> <span class="n">parent_selections</span> <span class="o">=</span> <span class="p">{</span><span class="n">SELECT_ALL</span><span class="p">}</span>
</span><span id="pushdown_projections-51"><a href="#pushdown_projections-51"><span class="linenos">51</span></a>
</span><span id="pushdown_projections-52"><a href="#pushdown_projections-52"><span class="linenos">52</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Union</span><span class="p">):</span>
</span><span id="pushdown_projections-53"><a href="#pushdown_projections-53"><span class="linenos">53</span></a> <span class="n">left</span><span class="p">,</span> <span class="n">right</span> <span class="o">=</span> <span class="n">scope</span><span class="o">.</span><span class="n">union_scopes</span>
</span><span id="pushdown_projections-54"><a href="#pushdown_projections-54"><span class="linenos">54</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">left</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="pushdown_projections-55"><a href="#pushdown_projections-55"><span class="linenos">55</span></a>
</span><span id="pushdown_projections-56"><a href="#pushdown_projections-56"><span class="linenos">56</span></a> <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="pushdown_projections-57"><a href="#pushdown_projections-57"><span class="linenos">57</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="n">parent_selections</span>
</span><span id="pushdown_projections-58"><a href="#pushdown_projections-58"><span class="linenos">58</span></a> <span class="k">elif</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">is_star</span> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">):</span>
</span><span id="pushdown_projections-59"><a href="#pushdown_projections-59"><span class="linenos">59</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">right</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span>
</span><span id="pushdown_projections-60"><a href="#pushdown_projections-60"><span class="linenos">60</span></a> <span class="n">right</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">alias_or_name</span>
</span><span id="pushdown_projections-61"><a href="#pushdown_projections-61"><span class="linenos">61</span></a> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">select</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="p">)</span>
</span><span id="pushdown_projections-62"><a href="#pushdown_projections-62"><span class="linenos">62</span></a> <span class="k">if</span> <span class="n">SELECT_ALL</span> <span class="ow">in</span> <span class="n">parent_selections</span> <span class="ow">or</span> <span class="n">select</span><span class="o">.</span><span class="n">alias_or_name</span> <span class="ow">in</span> <span class="n">parent_selections</span>
</span><span id="pushdown_projections-63"><a href="#pushdown_projections-63"><span class="linenos">63</span></a> <span class="p">]</span>
</span><span id="pushdown_projections-64"><a href="#pushdown_projections-64"><span class="linenos">64</span></a>
</span><span id="pushdown_projections-65"><a href="#pushdown_projections-65"><span class="linenos">65</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Select</span><span class="p">):</span>
</span><span id="pushdown_projections-66"><a href="#pushdown_projections-66"><span class="linenos">66</span></a> <span class="k">if</span> <span class="n">remove_unused_selections</span><span class="p">:</span>
</span><span id="pushdown_projections-67"><a href="#pushdown_projections-67"><span class="linenos">67</span></a> <span class="n">_remove_unused_selections</span><span class="p">(</span><span class="n">scope</span><span class="p">,</span> <span class="n">parent_selections</span><span class="p">,</span> <span class="n">schema</span><span class="p">,</span> <span class="n">alias_count</span><span class="p">)</span>
</span><span id="pushdown_projections-68"><a href="#pushdown_projections-68"><span class="linenos">68</span></a>
</span><span id="pushdown_projections-69"><a href="#pushdown_projections-69"><span class="linenos">69</span></a> <span class="k">if</span> <span class="n">scope</span><span class="o">.</span><span class="n">expression</span><span class="o">.</span><span class="n">is_star</span><span class="p">:</span>
</span><span id="pushdown_projections-70"><a href="#pushdown_projections-70"><span class="linenos">70</span></a> <span class="k">continue</span>
</span><span id="pushdown_projections-71"><a href="#pushdown_projections-71"><span class="linenos">71</span></a>
</span><span id="pushdown_projections-72"><a href="#pushdown_projections-72"><span class="linenos">72</span></a> <span class="c1"># Group columns by source name</span>
</span><span id="pushdown_projections-73"><a href="#pushdown_projections-73"><span class="linenos">73</span></a> <span class="n">selects</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>
</span><span id="pushdown_projections-74"><a href="#pushdown_projections-74"><span class="linenos">74</span></a> <span class="k">for</span> <span class="n">col</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">columns</span><span class="p">:</span>
</span><span id="pushdown_projections-75"><a href="#pushdown_projections-75"><span class="linenos">75</span></a> <span class="n">table_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">table</span>
</span><span id="pushdown_projections-76"><a href="#pushdown_projections-76"><span class="linenos">76</span></a> <span class="n">col_name</span> <span class="o">=</span> <span class="n">col</span><span class="o">.</span><span class="n">name</span>
</span><span id="pushdown_projections-77"><a href="#pushdown_projections-77"><span class="linenos">77</span></a> <span class="n">selects</span><span class="p">[</span><span class="n">table_name</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">col_name</span><span class="p">)</span>
</span><span id="pushdown_projections-78"><a href="#pushdown_projections-78"><span class="linenos">78</span></a>
</span><span id="pushdown_projections-79"><a href="#pushdown_projections-79"><span class="linenos">79</span></a> <span class="c1"># Push the selected columns down to the next scope</span>
</span><span id="pushdown_projections-80"><a href="#pushdown_projections-80"><span class="linenos">80</span></a> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">source</span><span class="p">)</span> <span class="ow">in</span> <span class="n">scope</span><span class="o">.</span><span class="n">selected_sources</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
</span><span id="pushdown_projections-81"><a href="#pushdown_projections-81"><span class="linenos">81</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">source</span><span class="p">,</span> <span class="n">Scope</span><span class="p">):</span>
</span><span id="pushdown_projections-82"><a href="#pushdown_projections-82"><span class="linenos">82</span></a> <span class="n">columns</span> <span class="o">=</span> <span class="n">selects</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">set</span><span class="p">()</span>
</span><span id="pushdown_projections-83"><a href="#pushdown_projections-83"><span class="linenos">83</span></a> <span class="n">referenced_columns</span><span class="p">[</span><span class="n">source</span><span class="p">]</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">columns</span><span class="p">)</span>
</span><span id="pushdown_projections-84"><a href="#pushdown_projections-84"><span class="linenos">84</span></a>
</span><span id="pushdown_projections-85"><a href="#pushdown_projections-85"><span class="linenos">85</span></a> <span class="n">column_aliases</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">alias_column_names</span>
</span><span id="pushdown_projections-86"><a href="#pushdown_projections-86"><span class="linenos">86</span></a> <span class="k">if</span> <span class="n">column_aliases</span><span class="p">:</span>
</span><span id="pushdown_projections-87"><a href="#pushdown_projections-87"><span class="linenos">87</span></a> <span class="n">source_column_alias_count</span><span class="p">[</span><span class="n">source</span><span class="p">]</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">column_aliases</span><span class="p">)</span>
</span><span id="pushdown_projections-88"><a href="#pushdown_projections-88"><span class="linenos">88</span></a>
</span><span id="pushdown_projections-89"><a href="#pushdown_projections-89"><span class="linenos">89</span></a> <span class="k">return</span> <span class="n">expression</span>
</span></pre></div>

File diff suppressed because it is too large Load diff

File diff suppressed because one or more lines are too long

File diff suppressed because it is too large Load diff

View file

@ -7824,7 +7824,7 @@
<div class="attr variable">
<span class="name">COMMANDS</span> =
<input id="Tokenizer.COMMANDS-view-value" class="view-value-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
<label class="view-value-button pdoc-button" for="Tokenizer.COMMANDS-view-value"></label><span class="default_value">{&lt;<a href="#TokenType.COMMAND">TokenType.COMMAND</a>: &#39;COMMAND&#39;&gt;, &lt;<a href="#TokenType.FETCH">TokenType.FETCH</a>: &#39;FETCH&#39;&gt;, &lt;<a href="#TokenType.SHOW">TokenType.SHOW</a>: &#39;SHOW&#39;&gt;, &lt;<a href="#TokenType.EXECUTE">TokenType.EXECUTE</a>: &#39;EXECUTE&#39;&gt;}</span>
<label class="view-value-button pdoc-button" for="Tokenizer.COMMANDS-view-value"></label><span class="default_value">{&lt;<a href="#TokenType.FETCH">TokenType.FETCH</a>: &#39;FETCH&#39;&gt;, &lt;<a href="#TokenType.SHOW">TokenType.SHOW</a>: &#39;SHOW&#39;&gt;, &lt;<a href="#TokenType.EXECUTE">TokenType.EXECUTE</a>: &#39;EXECUTE&#39;&gt;, &lt;<a href="#TokenType.COMMAND">TokenType.COMMAND</a>: &#39;COMMAND&#39;&gt;}</span>
</div>

View file

@ -286,9 +286,9 @@
</span><span id="L-189"><a href="#L-189"><span class="linenos">189</span></a>
</span><span id="L-190"><a href="#L-190"><span class="linenos">190</span></a> <span class="c1"># we use list here because expression.selects is mutated inside the loop</span>
</span><span id="L-191"><a href="#L-191"><span class="linenos">191</span></a> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="o">.</span><span class="n">copy</span><span class="p">():</span>
</span><span id="L-192"><a href="#L-192"><span class="linenos">192</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)</span>
</span><span id="L-192"><a href="#L-192"><span class="linenos">192</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">)</span>
</span><span id="L-193"><a href="#L-193"><span class="linenos">193</span></a>
</span><span id="L-194"><a href="#L-194"><span class="linenos">194</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">explode</span><span class="p">,</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)):</span>
</span><span id="L-194"><a href="#L-194"><span class="linenos">194</span></a> <span class="k">if</span> <span class="n">explode</span><span class="p">:</span>
</span><span id="L-195"><a href="#L-195"><span class="linenos">195</span></a> <span class="n">pos_alias</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
</span><span id="L-196"><a href="#L-196"><span class="linenos">196</span></a> <span class="n">explode_alias</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
</span><span id="L-197"><a href="#L-197"><span class="linenos">197</span></a>
@ -301,7 +301,7 @@
</span><span id="L-204"><a href="#L-204"><span class="linenos">204</span></a> <span class="n">alias</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">alias_</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">this</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">))</span>
</span><span id="L-205"><a href="#L-205"><span class="linenos">205</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="L-206"><a href="#L-206"><span class="linenos">206</span></a> <span class="n">alias</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">alias_</span><span class="p">(</span><span class="n">select</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">))</span>
</span><span id="L-207"><a href="#L-207"><span class="linenos">207</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">alias</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)</span>
</span><span id="L-207"><a href="#L-207"><span class="linenos">207</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">alias</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">)</span>
</span><span id="L-208"><a href="#L-208"><span class="linenos">208</span></a> <span class="k">assert</span> <span class="n">explode</span>
</span><span id="L-209"><a href="#L-209"><span class="linenos">209</span></a>
</span><span id="L-210"><a href="#L-210"><span class="linenos">210</span></a> <span class="n">is_posexplode</span> <span class="o">=</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)</span>
@ -844,9 +844,9 @@ other expressions. This transforms removes the precision from parameterized type
</span><span id="explode_to_unnest-190"><a href="#explode_to_unnest-190"><span class="linenos">190</span></a>
</span><span id="explode_to_unnest-191"><a href="#explode_to_unnest-191"><span class="linenos">191</span></a> <span class="c1"># we use list here because expression.selects is mutated inside the loop</span>
</span><span id="explode_to_unnest-192"><a href="#explode_to_unnest-192"><span class="linenos">192</span></a> <span class="k">for</span> <span class="n">select</span> <span class="ow">in</span> <span class="n">expression</span><span class="o">.</span><span class="n">selects</span><span class="o">.</span><span class="n">copy</span><span class="p">():</span>
</span><span id="explode_to_unnest-193"><a href="#explode_to_unnest-193"><span class="linenos">193</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)</span>
</span><span id="explode_to_unnest-193"><a href="#explode_to_unnest-193"><span class="linenos">193</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">)</span>
</span><span id="explode_to_unnest-194"><a href="#explode_to_unnest-194"><span class="linenos">194</span></a>
</span><span id="explode_to_unnest-195"><a href="#explode_to_unnest-195"><span class="linenos">195</span></a> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">explode</span><span class="p">,</span> <span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)):</span>
</span><span id="explode_to_unnest-195"><a href="#explode_to_unnest-195"><span class="linenos">195</span></a> <span class="k">if</span> <span class="n">explode</span><span class="p">:</span>
</span><span id="explode_to_unnest-196"><a href="#explode_to_unnest-196"><span class="linenos">196</span></a> <span class="n">pos_alias</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
</span><span id="explode_to_unnest-197"><a href="#explode_to_unnest-197"><span class="linenos">197</span></a> <span class="n">explode_alias</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
</span><span id="explode_to_unnest-198"><a href="#explode_to_unnest-198"><span class="linenos">198</span></a>
@ -859,7 +859,7 @@ other expressions. This transforms removes the precision from parameterized type
</span><span id="explode_to_unnest-205"><a href="#explode_to_unnest-205"><span class="linenos">205</span></a> <span class="n">alias</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">alias_</span><span class="p">(</span><span class="n">select</span><span class="o">.</span><span class="n">this</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="n">copy</span><span class="o">=</span><span class="kc">False</span><span class="p">))</span>
</span><span id="explode_to_unnest-206"><a href="#explode_to_unnest-206"><span class="linenos">206</span></a> <span class="k">else</span><span class="p">:</span>
</span><span id="explode_to_unnest-207"><a href="#explode_to_unnest-207"><span class="linenos">207</span></a> <span class="n">alias</span> <span class="o">=</span> <span class="n">select</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">alias_</span><span class="p">(</span><span class="n">select</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">))</span>
</span><span id="explode_to_unnest-208"><a href="#explode_to_unnest-208"><span class="linenos">208</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">alias</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)</span>
</span><span id="explode_to_unnest-208"><a href="#explode_to_unnest-208"><span class="linenos">208</span></a> <span class="n">explode</span> <span class="o">=</span> <span class="n">alias</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">exp</span><span class="o">.</span><span class="n">Explode</span><span class="p">)</span>
</span><span id="explode_to_unnest-209"><a href="#explode_to_unnest-209"><span class="linenos">209</span></a> <span class="k">assert</span> <span class="n">explode</span>
</span><span id="explode_to_unnest-210"><a href="#explode_to_unnest-210"><span class="linenos">210</span></a>
</span><span id="explode_to_unnest-211"><a href="#explode_to_unnest-211"><span class="linenos">211</span></a> <span class="n">is_posexplode</span> <span class="o">=</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">explode</span><span class="p">,</span> <span class="n">exp</span><span class="o">.</span><span class="n">Posexplode</span><span class="p">)</span>

View file

@ -1019,11 +1019,11 @@ def posexplode(col: ColumnOrName) -> Column:
def explode_outer(col: ColumnOrName) -> Column:
return Column.invoke_anonymous_function(col, "EXPLODE_OUTER")
return Column.invoke_expression_over_column(col, expression.ExplodeOuter)
def posexplode_outer(col: ColumnOrName) -> Column:
return Column.invoke_anonymous_function(col, "POSEXPLODE_OUTER")
return Column.invoke_expression_over_column(col, expression.PosexplodeOuter)
def get_json_object(col: ColumnOrName, path: str) -> Column:

View file

@ -10,6 +10,7 @@ from sqlglot.tokens import TokenType
class Databricks(Spark):
class Parser(Spark.Parser):
LOG_DEFAULTS_TO_LN = True
STRICT_CAST = True
FUNCTIONS = {
**Spark.Parser.FUNCTIONS,
@ -51,6 +52,8 @@ class Databricks(Spark):
exp.ToChar: lambda self, e: self.function_fallback_sql(e),
}
TRANSFORMS.pop(exp.TryCast)
def columndef_sql(self, expression: exp.ColumnDef, sep: str = " ") -> str:
constraint = expression.find(exp.GeneratedAsIdentityColumnConstraint)
kind = expression.args.get("kind")

View file

@ -133,6 +133,10 @@ class DuckDB(Dialect):
"UINTEGER": TokenType.UINT,
"USMALLINT": TokenType.USMALLINT,
"UTINYINT": TokenType.UTINYINT,
"TIMESTAMP_S": TokenType.TIMESTAMP_S,
"TIMESTAMP_MS": TokenType.TIMESTAMP_MS,
"TIMESTAMP_NS": TokenType.TIMESTAMP_NS,
"TIMESTAMP_US": TokenType.TIMESTAMP,
}
class Parser(parser.Parser):
@ -321,6 +325,9 @@ class DuckDB(Dialect):
exp.DataType.Type.UINT: "UINTEGER",
exp.DataType.Type.VARBINARY: "BLOB",
exp.DataType.Type.VARCHAR: "TEXT",
exp.DataType.Type.TIMESTAMP_S: "TIMESTAMP_S",
exp.DataType.Type.TIMESTAMP_MS: "TIMESTAMP_MS",
exp.DataType.Type.TIMESTAMP_NS: "TIMESTAMP_NS",
}
STAR_MAPPING = {**generator.Generator.STAR_MAPPING, "except": "EXCLUDE"}

View file

@ -82,7 +82,6 @@ class Oracle(Dialect):
this=self._parse_format_json(self._parse_bitwise()),
order=self._parse_order(),
),
"JSON_TABLE": lambda self: self._parse_json_table(),
"XMLTABLE": _parse_xml_table,
}
@ -96,29 +95,6 @@ class Oracle(Dialect):
# Reference: https://stackoverflow.com/a/336455
DISTINCT_TOKENS = {TokenType.DISTINCT, TokenType.UNIQUE}
# Note: this is currently incomplete; it only implements the "JSON_value_column" part
def _parse_json_column_def(self) -> exp.JSONColumnDef:
this = self._parse_id_var()
kind = self._parse_types(allow_identifiers=False)
path = self._match_text_seq("PATH") and self._parse_string()
return self.expression(exp.JSONColumnDef, this=this, kind=kind, path=path)
def _parse_json_table(self) -> exp.JSONTable:
this = self._parse_format_json(self._parse_bitwise())
path = self._match(TokenType.COMMA) and self._parse_string()
error_handling = self._parse_on_handling("ERROR", "ERROR", "NULL")
empty_handling = self._parse_on_handling("EMPTY", "ERROR", "NULL")
self._match(TokenType.COLUMN)
expressions = self._parse_wrapped_csv(self._parse_json_column_def, optional=True)
return exp.JSONTable(
this=this,
expressions=expressions,
path=path,
error_handling=error_handling,
empty_handling=empty_handling,
)
def _parse_json_array(self, expr_type: t.Type[E], **kwargs) -> E:
return self.expression(
expr_type,

View file

@ -34,7 +34,7 @@ def _approx_distinct_sql(self: Presto.Generator, expression: exp.ApproxDistinct)
def _explode_to_unnest_sql(self: Presto.Generator, expression: exp.Lateral) -> str:
if isinstance(expression.this, (exp.Explode, exp.Posexplode)):
if isinstance(expression.this, exp.Explode):
expression = expression.copy()
return self.sql(
exp.Join(

View file

@ -58,6 +58,11 @@ class Redshift(Postgres):
"STRTOL": exp.FromBase.from_arg_list,
}
NO_PAREN_FUNCTION_PARSERS = {
**Postgres.Parser.NO_PAREN_FUNCTION_PARSERS,
"APPROXIMATE": lambda self: self._parse_approximate_count(),
}
def _parse_table(
self,
schema: bool = False,
@ -93,11 +98,22 @@ class Redshift(Postgres):
return this
def _parse_convert(self, strict: bool) -> t.Optional[exp.Expression]:
def _parse_convert(
self, strict: bool, safe: t.Optional[bool] = None
) -> t.Optional[exp.Expression]:
to = self._parse_types()
self._match(TokenType.COMMA)
this = self._parse_bitwise()
return self.expression(exp.TryCast, this=this, to=to)
return self.expression(exp.TryCast, this=this, to=to, safe=safe)
def _parse_approximate_count(self) -> t.Optional[exp.ApproxDistinct]:
index = self._index - 1
func = self._parse_function()
if isinstance(func, exp.Count) and isinstance(func.this, exp.Distinct):
return self.expression(exp.ApproxDistinct, this=seq_get(func.this.expressions, 0))
self._retreat(index)
return None
class Tokenizer(Postgres.Tokenizer):
BIT_STRINGS = []
@ -144,6 +160,7 @@ class Redshift(Postgres):
**Postgres.Generator.TRANSFORMS,
exp.Concat: concat_to_dpipe_sql,
exp.ConcatWs: concat_ws_to_dpipe_sql,
exp.ApproxDistinct: lambda self, e: f"APPROXIMATE COUNT(DISTINCT {self.sql(e, 'this')})",
exp.CurrentTimestamp: lambda self, e: "SYSDATE",
exp.DateAdd: lambda self, e: self.func(
"DATEADD", exp.var(e.text("unit") or "day"), e.expression, e.this

View file

@ -76,6 +76,9 @@ class Spark(Spark2):
exp.TimestampAdd: lambda self, e: self.func(
"DATEADD", e.args.get("unit") or "DAY", e.expression, e.this
),
exp.TryCast: lambda self, e: self.trycast_sql(e)
if e.args.get("safe")
else self.cast_sql(e),
}
TRANSFORMS.pop(exp.AnyValue)
TRANSFORMS.pop(exp.DateDiff)

View file

@ -477,7 +477,9 @@ class TSQL(Dialect):
returns.set("table", table)
return returns
def _parse_convert(self, strict: bool) -> t.Optional[exp.Expression]:
def _parse_convert(
self, strict: bool, safe: t.Optional[bool] = None
) -> t.Optional[exp.Expression]:
to = self._parse_types()
self._match(TokenType.COMMA)
this = self._parse_conjunction()
@ -513,12 +515,13 @@ class TSQL(Dialect):
exp.Cast if strict else exp.TryCast,
to=to,
this=self.expression(exp.TimeToStr, this=this, format=format_norm),
safe=safe,
)
elif to.this == DataType.Type.TEXT:
return self.expression(exp.TimeToStr, this=this, format=format_norm)
# Entails a simple cast without any format requirement
return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to)
return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to, safe=safe)
def _parse_user_defined_function(
self, kind: t.Optional[TokenType] = None

View file

@ -105,7 +105,7 @@ class RowReader:
return self.row[self.columns[column]]
class Tables(AbstractMappingSchema[Table]):
class Tables(AbstractMappingSchema):
pass

View file

@ -487,7 +487,7 @@ class Expression(metaclass=_Expression):
"""
for node, _, _ in self.dfs(prune=lambda n, p, *_: p and not type(n) is self.__class__):
if not type(node) is self.__class__:
yield node.unnest() if unnest else node
yield node.unnest() if unnest and not isinstance(node, Subquery) else node
def __str__(self) -> str:
return self.sql()
@ -2107,7 +2107,7 @@ class LockingProperty(Property):
arg_types = {
"this": False,
"kind": True,
"for_or_in": True,
"for_or_in": False,
"lock_type": True,
"override": False,
}
@ -3605,6 +3605,9 @@ class DataType(Expression):
TIMESTAMP = auto()
TIMESTAMPLTZ = auto()
TIMESTAMPTZ = auto()
TIMESTAMP_S = auto()
TIMESTAMP_MS = auto()
TIMESTAMP_NS = auto()
TINYINT = auto()
TSMULTIRANGE = auto()
TSRANGE = auto()
@ -3661,6 +3664,9 @@ class DataType(Expression):
Type.TIMESTAMP,
Type.TIMESTAMPTZ,
Type.TIMESTAMPLTZ,
Type.TIMESTAMP_S,
Type.TIMESTAMP_MS,
Type.TIMESTAMP_NS,
Type.DATE,
Type.DATETIME,
Type.DATETIME64,
@ -4286,7 +4292,7 @@ class Case(Func):
class Cast(Func):
arg_types = {"this": True, "to": True, "format": False}
arg_types = {"this": True, "to": True, "format": False, "safe": False}
@property
def name(self) -> str:
@ -4538,6 +4544,18 @@ class Explode(Func):
pass
class ExplodeOuter(Explode):
pass
class Posexplode(Explode):
pass
class PosexplodeOuter(Posexplode):
pass
class Floor(Func):
arg_types = {"this": True, "decimals": False}
@ -4621,14 +4639,18 @@ class JSONArrayAgg(Func):
# https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/JSON_TABLE.html
# Note: parsing of JSON column definitions is currently incomplete.
class JSONColumnDef(Expression):
arg_types = {"this": True, "kind": False, "path": False}
arg_types = {"this": False, "kind": False, "path": False, "nested_schema": False}
class JSONSchema(Expression):
arg_types = {"expressions": True}
# # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/JSON_TABLE.html
class JSONTable(Func):
arg_types = {
"this": True,
"expressions": True,
"schema": True,
"path": False,
"error_handling": False,
"empty_handling": False,
@ -4790,10 +4812,6 @@ class Nvl2(Func):
arg_types = {"this": True, "true": True, "false": False}
class Posexplode(Func):
pass
# https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-predict#mlpredict_function
class Predict(Func):
arg_types = {"this": True, "expression": True, "params_struct": False}

View file

@ -1226,9 +1226,10 @@ class Generator:
kind = expression.args.get("kind")
this = f" {self.sql(expression, 'this')}" if expression.this else ""
for_or_in = expression.args.get("for_or_in")
for_or_in = f" {for_or_in}" if for_or_in else ""
lock_type = expression.args.get("lock_type")
override = " OVERRIDE" if expression.args.get("override") else ""
return f"LOCKING {kind}{this} {for_or_in} {lock_type}{override}"
return f"LOCKING {kind}{this}{for_or_in} {lock_type}{override}"
def withdataproperty_sql(self, expression: exp.WithDataProperty) -> str:
data_sql = f"WITH {'NO ' if expression.args.get('no') else ''}DATA"
@ -2179,13 +2180,21 @@ class Generator:
)
def jsoncolumndef_sql(self, expression: exp.JSONColumnDef) -> str:
path = self.sql(expression, "path")
path = f" PATH {path}" if path else ""
nested_schema = self.sql(expression, "nested_schema")
if nested_schema:
return f"NESTED{path} {nested_schema}"
this = self.sql(expression, "this")
kind = self.sql(expression, "kind")
kind = f" {kind}" if kind else ""
path = self.sql(expression, "path")
path = f" PATH {path}" if path else ""
return f"{this}{kind}{path}"
def jsonschema_sql(self, expression: exp.JSONSchema) -> str:
return self.func("COLUMNS", *expression.expressions)
def jsontable_sql(self, expression: exp.JSONTable) -> str:
this = self.sql(expression, "this")
path = self.sql(expression, "path")
@ -2194,9 +2203,9 @@ class Generator:
error_handling = f" {error_handling}" if error_handling else ""
empty_handling = expression.args.get("empty_handling")
empty_handling = f" {empty_handling}" if empty_handling else ""
columns = f" COLUMNS ({self.expressions(expression, skip_first=True)})"
schema = self.sql(expression, "schema")
return self.func(
"JSON_TABLE", this, suffix=f"{path}{error_handling}{empty_handling}{columns})"
"JSON_TABLE", this, suffix=f"{path}{error_handling}{empty_handling} {schema})"
)
def openjsoncolumndef_sql(self, expression: exp.OpenJSONColumnDef) -> str:

View file

@ -441,6 +441,14 @@ def first(it: t.Iterable[T]) -> T:
def merge_ranges(ranges: t.List[t.Tuple[A, A]]) -> t.List[t.Tuple[A, A]]:
"""
Merges a sequence of ranges, represented as tuples (low, high) whose values
belong to some totally-ordered set.
Example:
>>> merge_ranges([(1, 3), (2, 6)])
[(1, 6)]
"""
if not ranges:
return []

View file

@ -6,6 +6,7 @@ from sqlglot import exp
from sqlglot.errors import OptimizeError
from sqlglot.generator import cached_generator
from sqlglot.helper import while_changing
from sqlglot.optimizer.scope import find_all_in_scope
from sqlglot.optimizer.simplify import flatten, rewrite_between, uniq_sort
logger = logging.getLogger("sqlglot")
@ -63,15 +64,33 @@ def normalize(expression: exp.Expression, dnf: bool = False, max_distance: int =
return expression
def normalized(expression, dnf=False):
ancestor, root = (exp.And, exp.Or) if dnf else (exp.Or, exp.And)
return not any(connector.find_ancestor(ancestor) for connector in expression.find_all(root))
def normalization_distance(expression, dnf=False):
def normalized(expression: exp.Expression, dnf: bool = False) -> bool:
"""
The difference in the number of predicates between the current expression and the normalized form.
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 or not 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).
"""
ancestor, root = (exp.And, exp.Or) if dnf else (exp.Or, exp.And)
return not any(
connector.find_ancestor(ancestor) for connector in find_all_in_scope(expression, root)
)
def normalization_distance(expression: exp.Expression, dnf: bool = False) -> int:
"""
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.
@ -82,10 +101,12 @@ def normalization_distance(expression, dnf=False):
4
Args:
expression (sqlglot.Expression): expression to compute distance
dnf (bool): compute to dnf distance instead
expression: The expression to compute the normalization distance for.
dnf: Whether or not 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).
Returns:
int: difference
The normalization distance.
"""
return sum(_predicate_lengths(expression, dnf)) - (
sum(1 for _ in expression.find_all(exp.Connector)) + 1

View file

@ -39,10 +39,14 @@ def optimize_joins(expression):
if len(other_table_names(dep)) < 2:
continue
operator = type(on)
for predicate in on.flatten():
if name in exp.column_table_names(predicate):
predicate.replace(exp.true())
join.on(predicate, copy=False)
predicate = exp._combine(
[join.args.get("on"), predicate], operator, copy=False
)
join.on(predicate, append=False, copy=False)
expression = reorder_joins(expression)
expression = normalize(expression)

View file

@ -9,7 +9,9 @@ from sqlglot.schema import ensure_schema
SELECT_ALL = object()
# Selection to use if selection list is empty
DEFAULT_SELECTION = lambda: alias("1", "_")
DEFAULT_SELECTION = lambda is_agg: alias(
exp.Max(this=exp.Literal.number(1)) if is_agg else "1", "_"
)
def pushdown_projections(expression, schema=None, remove_unused_selections=True):
@ -98,6 +100,7 @@ def _remove_unused_selections(scope, parent_selections, schema, alias_count):
new_selections = []
removed = False
star = False
is_agg = False
select_all = SELECT_ALL in parent_selections
@ -112,6 +115,9 @@ def _remove_unused_selections(scope, parent_selections, schema, alias_count):
star = True
removed = True
if not is_agg and selection.find(exp.AggFunc):
is_agg = True
if star:
resolver = Resolver(scope, schema)
names = {s.alias_or_name for s in new_selections}
@ -124,7 +130,7 @@ def _remove_unused_selections(scope, parent_selections, schema, alias_count):
# If there are no remaining selections, just select a single constant
if not new_selections:
new_selections.append(DEFAULT_SELECTION())
new_selections.append(DEFAULT_SELECTION(is_agg))
scope.expression.select(*new_selections, append=False, copy=False)

View file

@ -137,8 +137,8 @@ class Scope:
if not self._collected:
self._collect()
def walk(self, bfs=True):
return walk_in_scope(self.expression, bfs=bfs)
def walk(self, bfs=True, prune=None):
return walk_in_scope(self.expression, bfs=bfs, prune=None)
def find(self, *expression_types, bfs=True):
return find_in_scope(self.expression, expression_types, bfs=bfs)
@ -731,7 +731,7 @@ def _traverse_ddl(scope):
yield from _traverse_scope(query_scope)
def walk_in_scope(expression, bfs=True):
def walk_in_scope(expression, bfs=True, prune=None):
"""
Returns a generator object which visits all nodes in the syntrax tree, stopping at
nodes that start child scopes.
@ -740,16 +740,20 @@ def walk_in_scope(expression, bfs=True):
expression (exp.Expression):
bfs (bool): if set to True the BFS traversal order will be applied,
otherwise the DFS traversal will be used instead.
prune ((node, parent, arg_key) -> bool): callable that returns True if
the generator should stop traversing this branch of the tree.
Yields:
tuple[exp.Expression, Optional[exp.Expression], str]: node, parent, arg key
"""
# We'll use this variable to pass state into the dfs generator.
# Whenever we set it to True, we exclude a subtree from traversal.
prune = False
crossed_scope_boundary = False
for node, parent, key in expression.walk(bfs=bfs, prune=lambda *_: prune):
prune = False
for node, parent, key in expression.walk(
bfs=bfs, prune=lambda *args: crossed_scope_boundary or (prune and prune(*args))
):
crossed_scope_boundary = False
yield node, parent, key
@ -765,7 +769,7 @@ def walk_in_scope(expression, bfs=True):
or isinstance(node, exp.UDTF)
or isinstance(node, exp.Subqueryable)
):
prune = True
crossed_scope_boundary = True
if isinstance(node, (exp.Subquery, exp.UDTF)):
# The following args are not actually in the inner scope, so we should visit them

View file

@ -5,9 +5,11 @@ import typing as t
from collections import deque
from decimal import Decimal
import sqlglot
from sqlglot import exp
from sqlglot.generator import cached_generator
from sqlglot.helper import first, merge_ranges, while_changing
from sqlglot.optimizer.scope import find_all_in_scope, walk_in_scope
# Final means that an expression should not be simplified
FINAL = "final"
@ -17,7 +19,7 @@ class UnsupportedUnit(Exception):
pass
def simplify(expression):
def simplify(expression, constant_propagation=False):
"""
Rewrite sqlglot AST to simplify expressions.
@ -29,6 +31,8 @@ def simplify(expression):
Args:
expression (sqlglot.Expression): expression to simplify
constant_propagation: whether or not the constant propagation rule should be used
Returns:
sqlglot.Expression: simplified expression
"""
@ -67,13 +71,16 @@ def simplify(expression):
node = absorb_and_eliminate(node, root)
node = simplify_concat(node)
if constant_propagation:
node = propagate_constants(node, root)
exp.replace_children(node, lambda e: _simplify(e, False))
# Post-order transformations
node = simplify_not(node)
node = flatten(node)
node = simplify_connectors(node, root)
node = remove_compliments(node, root)
node = remove_complements(node, root)
node = simplify_coalesce(node)
node.parent = expression.parent
node = simplify_literals(node, root)
@ -287,19 +294,19 @@ def _simplify_comparison(expression, left, right, or_=False):
return None
def remove_compliments(expression, root=True):
def remove_complements(expression, root=True):
"""
Removing compliments.
Removing complements.
A AND NOT A -> FALSE
A OR NOT A -> TRUE
"""
if isinstance(expression, exp.Connector) and (root or not expression.same_parent):
compliment = exp.false() if isinstance(expression, exp.And) else exp.true()
complement = exp.false() if isinstance(expression, exp.And) else exp.true()
for a, b in itertools.permutations(expression.flatten(), 2):
if is_complement(a, b):
return compliment
return complement
return expression
@ -369,6 +376,51 @@ def absorb_and_eliminate(expression, root=True):
return expression
def propagate_constants(expression, root=True):
"""
Propagate constants for conjunctions in DNF:
SELECT * FROM t WHERE a = b AND b = 5 becomes
SELECT * FROM t WHERE a = 5 AND b = 5
Reference: https://www.sqlite.org/optoverview.html
"""
if (
isinstance(expression, exp.And)
and (root or not expression.same_parent)
and sqlglot.optimizer.normalize.normalized(expression, dnf=True)
):
constant_mapping = {}
for expr, *_ in walk_in_scope(expression, prune=lambda node, *_: isinstance(node, exp.If)):
if isinstance(expr, exp.EQ):
l, r = expr.left, expr.right
# TODO: create a helper that can be used to detect nested literal expressions such
# as CAST(123456 AS BIGINT), since we usually want to treat those as literals too
if isinstance(l, exp.Column) and isinstance(r, exp.Literal):
pass
elif isinstance(r, exp.Column) and isinstance(l, exp.Literal):
l, r = r, l
else:
continue
constant_mapping[l] = (id(l), r)
if constant_mapping:
for column in find_all_in_scope(expression, exp.Column):
parent = column.parent
column_id, constant = constant_mapping.get(column) or (None, None)
if (
column_id is not None
and id(column) != column_id
and not (isinstance(parent, exp.Is) and isinstance(parent.expression, exp.Null))
):
column.replace(constant.copy())
return expression
INVERSE_DATE_OPS: t.Dict[t.Type[exp.Expression], t.Type[exp.Expression]] = {
exp.DateAdd: exp.Sub,
exp.DateSub: exp.Add,
@ -609,21 +661,38 @@ SAFE_CONCATS = (exp.SafeConcat, exp.SafeDPipe)
def simplify_concat(expression):
"""Reduces all groups that contain string literals by concatenating them."""
if not isinstance(expression, CONCATS) or isinstance(expression, exp.ConcatWs):
if not isinstance(expression, CONCATS) or (
# We can't reduce a CONCAT_WS call if we don't statically know the separator
isinstance(expression, exp.ConcatWs)
and not expression.expressions[0].is_string
):
return expression
if isinstance(expression, exp.ConcatWs):
sep_expr, *expressions = expression.expressions
sep = sep_expr.name
concat_type = exp.ConcatWs
else:
expressions = expression.expressions
sep = ""
concat_type = exp.SafeConcat if isinstance(expression, SAFE_CONCATS) else exp.Concat
new_args = []
for is_string_group, group in itertools.groupby(
expression.expressions or expression.flatten(), lambda e: e.is_string
expressions or expression.flatten(), lambda e: e.is_string
):
if is_string_group:
new_args.append(exp.Literal.string("".join(string.name for string in group)))
new_args.append(exp.Literal.string(sep.join(string.name for string in group)))
else:
new_args.extend(group)
# Ensures we preserve the right concat type, i.e. whether it's "safe" or not
concat_type = exp.SafeConcat if isinstance(expression, SAFE_CONCATS) else exp.Concat
return new_args[0] if len(new_args) == 1 else concat_type(expressions=new_args)
if len(new_args) == 1 and new_args[0].is_string:
return new_args[0]
if concat_type is exp.ConcatWs:
new_args = [sep_expr] + new_args
return concat_type(expressions=new_args)
DateRange = t.Tuple[datetime.date, datetime.date]

View file

@ -160,6 +160,9 @@ class Parser(metaclass=_Parser):
TokenType.TIME,
TokenType.TIMETZ,
TokenType.TIMESTAMP,
TokenType.TIMESTAMP_S,
TokenType.TIMESTAMP_MS,
TokenType.TIMESTAMP_NS,
TokenType.TIMESTAMPTZ,
TokenType.TIMESTAMPLTZ,
TokenType.DATETIME,
@ -792,17 +795,18 @@ class Parser(metaclass=_Parser):
"DECODE": lambda self: self._parse_decode(),
"EXTRACT": lambda self: self._parse_extract(),
"JSON_OBJECT": lambda self: self._parse_json_object(),
"JSON_TABLE": lambda self: self._parse_json_table(),
"LOG": lambda self: self._parse_logarithm(),
"MATCH": lambda self: self._parse_match_against(),
"OPENJSON": lambda self: self._parse_open_json(),
"POSITION": lambda self: self._parse_position(),
"PREDICT": lambda self: self._parse_predict(),
"SAFE_CAST": lambda self: self._parse_cast(False),
"SAFE_CAST": lambda self: self._parse_cast(False, safe=True),
"STRING_AGG": lambda self: self._parse_string_agg(),
"SUBSTRING": lambda self: self._parse_substring(),
"TRIM": lambda self: self._parse_trim(),
"TRY_CAST": lambda self: self._parse_cast(False),
"TRY_CONVERT": lambda self: self._parse_convert(False),
"TRY_CAST": lambda self: self._parse_cast(False, safe=True),
"TRY_CONVERT": lambda self: self._parse_convert(False, safe=True),
}
QUERY_MODIFIER_PARSERS = {
@ -4135,7 +4139,7 @@ class Parser(metaclass=_Parser):
return self.expression(exp.AnyValue, this=this, having=having, max=is_max)
def _parse_cast(self, strict: bool) -> exp.Expression:
def _parse_cast(self, strict: bool, safe: t.Optional[bool] = None) -> exp.Expression:
this = self._parse_conjunction()
if not self._match(TokenType.ALIAS):
@ -4176,7 +4180,9 @@ class Parser(metaclass=_Parser):
return this
return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to, format=fmt)
return self.expression(
exp.Cast if strict else exp.TryCast, this=this, to=to, format=fmt, safe=safe
)
def _parse_concat(self) -> t.Optional[exp.Expression]:
args = self._parse_csv(self._parse_conjunction)
@ -4230,7 +4236,9 @@ class Parser(metaclass=_Parser):
order = self._parse_order(this=seq_get(args, 0))
return self.expression(exp.GroupConcat, this=order, separator=seq_get(args, 1))
def _parse_convert(self, strict: bool) -> t.Optional[exp.Expression]:
def _parse_convert(
self, strict: bool, safe: t.Optional[bool] = None
) -> t.Optional[exp.Expression]:
this = self._parse_bitwise()
if self._match(TokenType.USING):
@ -4242,7 +4250,7 @@ class Parser(metaclass=_Parser):
else:
to = None
return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to)
return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to, safe=safe)
def _parse_decode(self) -> t.Optional[exp.Decode | exp.Case]:
"""
@ -4347,6 +4355,50 @@ class Parser(metaclass=_Parser):
encoding=encoding,
)
# Note: this is currently incomplete; it only implements the "JSON_value_column" part
def _parse_json_column_def(self) -> exp.JSONColumnDef:
if not self._match_text_seq("NESTED"):
this = self._parse_id_var()
kind = self._parse_types(allow_identifiers=False)
nested = None
else:
this = None
kind = None
nested = True
path = self._match_text_seq("PATH") and self._parse_string()
nested_schema = nested and self._parse_json_schema()
return self.expression(
exp.JSONColumnDef,
this=this,
kind=kind,
path=path,
nested_schema=nested_schema,
)
def _parse_json_schema(self) -> exp.JSONSchema:
self._match_text_seq("COLUMNS")
return self.expression(
exp.JSONSchema,
expressions=self._parse_wrapped_csv(self._parse_json_column_def, optional=True),
)
def _parse_json_table(self) -> exp.JSONTable:
this = self._parse_format_json(self._parse_bitwise())
path = self._match(TokenType.COMMA) and self._parse_string()
error_handling = self._parse_on_handling("ERROR", "ERROR", "NULL")
empty_handling = self._parse_on_handling("EMPTY", "ERROR", "NULL")
schema = self._parse_json_schema()
return exp.JSONTable(
this=this,
schema=schema,
path=path,
error_handling=error_handling,
empty_handling=empty_handling,
)
def _parse_logarithm(self) -> exp.Func:
# Default argument order is base, expression
args = self._parse_csv(self._parse_range)
@ -4973,7 +5025,17 @@ class Parser(metaclass=_Parser):
self._match(TokenType.ON)
on = self._parse_conjunction()
return self.expression(
exp.Merge,
this=target,
using=using,
on=on,
expressions=self._parse_when_matched(),
)
def _parse_when_matched(self) -> t.List[exp.When]:
whens = []
while self._match(TokenType.WHEN):
matched = not self._match(TokenType.NOT)
self._match_text_seq("MATCHED")
@ -5020,14 +5082,7 @@ class Parser(metaclass=_Parser):
then=then,
)
)
return self.expression(
exp.Merge,
this=target,
using=using,
on=on,
expressions=whens,
)
return whens
def _parse_show(self) -> t.Optional[exp.Expression]:
parser = self._find_parser(self.SHOW_PARSERS, self.SHOW_TRIE)

View file

@ -5,7 +5,6 @@ import typing as t
import sqlglot
from sqlglot import expressions as exp
from sqlglot._typing import T
from sqlglot.dialects.dialect import Dialect
from sqlglot.errors import ParseError, SchemaError
from sqlglot.helper import dict_depth
@ -71,7 +70,7 @@ class Schema(abc.ABC):
def get_column_type(
self,
table: exp.Table | str,
column: exp.Column,
column: exp.Column | str,
dialect: DialectType = None,
normalize: t.Optional[bool] = None,
) -> exp.DataType:
@ -88,6 +87,28 @@ class Schema(abc.ABC):
The resulting column type.
"""
def has_column(
self,
table: exp.Table | str,
column: exp.Column | str,
dialect: DialectType = None,
normalize: t.Optional[bool] = None,
) -> bool:
"""
Returns whether or not `column` appears in `table`'s schema.
Args:
table: the source table.
column: the target column.
dialect: the SQL dialect that will be used to parse `table` if it's a string.
normalize: whether to normalize identifiers according to the dialect of interest.
Returns:
True if the column appears in the schema, False otherwise.
"""
name = column if isinstance(column, str) else column.name
return name in self.column_names(table, dialect=dialect, normalize=normalize)
@property
@abc.abstractmethod
def supported_table_args(self) -> t.Tuple[str, ...]:
@ -101,7 +122,7 @@ class Schema(abc.ABC):
return True
class AbstractMappingSchema(t.Generic[T]):
class AbstractMappingSchema:
def __init__(
self,
mapping: t.Optional[t.Dict] = None,
@ -140,7 +161,7 @@ class AbstractMappingSchema(t.Generic[T]):
def find(
self, table: exp.Table, trie: t.Optional[t.Dict] = None, raise_on_missing: bool = True
) -> t.Optional[T]:
) -> t.Optional[t.Any]:
parts = self.table_parts(table)[0 : len(self.supported_table_args)]
value, trie = in_trie(self.mapping_trie if trie is None else trie, parts)
@ -170,7 +191,7 @@ class AbstractMappingSchema(t.Generic[T]):
)
class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
class MappingSchema(AbstractMappingSchema, Schema):
"""
Schema based on a nested mapping.
@ -287,7 +308,7 @@ class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
def get_column_type(
self,
table: exp.Table | str,
column: exp.Column,
column: exp.Column | str,
dialect: DialectType = None,
normalize: t.Optional[bool] = None,
) -> exp.DataType:
@ -304,10 +325,26 @@ class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
if isinstance(column_type, exp.DataType):
return column_type
elif isinstance(column_type, str):
return self._to_data_type(column_type.upper(), dialect=dialect)
return self._to_data_type(column_type, dialect=dialect)
return exp.DataType.build("unknown")
def has_column(
self,
table: exp.Table | str,
column: exp.Column | str,
dialect: DialectType = None,
normalize: t.Optional[bool] = None,
) -> bool:
normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
normalized_column_name = self._normalize_name(
column if isinstance(column, str) else column.this, dialect=dialect, normalize=normalize
)
table_schema = self.find(normalized_table, raise_on_missing=False)
return normalized_column_name in table_schema if table_schema else False
def _normalize(self, schema: t.Dict) -> t.Dict:
"""
Normalizes all identifiers in the schema.

View file

@ -121,6 +121,9 @@ class TokenType(AutoName):
TIMESTAMP = auto()
TIMESTAMPTZ = auto()
TIMESTAMPLTZ = auto()
TIMESTAMP_S = auto()
TIMESTAMP_MS = auto()
TIMESTAMP_NS = auto()
DATETIME = auto()
DATETIME64 = auto()
DATE = auto()

View file

@ -189,9 +189,9 @@ def explode_to_unnest(index_offset: int = 0) -> t.Callable[[exp.Expression], exp
# we use list here because expression.selects is mutated inside the loop
for select in expression.selects.copy():
explode = select.find(exp.Explode, exp.Posexplode)
explode = select.find(exp.Explode)
if isinstance(explode, (exp.Explode, exp.Posexplode)):
if explode:
pos_alias = ""
explode_alias = ""
@ -204,7 +204,7 @@ def explode_to_unnest(index_offset: int = 0) -> t.Callable[[exp.Expression], exp
alias = select.replace(exp.alias_(select.this, "", copy=False))
else:
alias = select.replace(exp.alias_(select, ""))
explode = alias.find(exp.Explode, exp.Posexplode)
explode = alias.find(exp.Explode)
assert explode
is_posexplode = isinstance(explode, exp.Posexplode)

View file

@ -790,3 +790,11 @@ class TestDuckDB(Validator):
"duckdb": "ALTER TABLE db.t1 RENAME TO t2",
},
)
def test_timestamps_with_units(self):
self.validate_all(
"SELECT w::TIMESTAMP_S, x::TIMESTAMP_MS, y::TIMESTAMP_US, z::TIMESTAMP_NS",
write={
"duckdb": "SELECT CAST(w AS TIMESTAMP_S), CAST(x AS TIMESTAMP_MS), CAST(y AS TIMESTAMP), CAST(z AS TIMESTAMP_NS)",
},
)

View file

@ -65,6 +65,9 @@ class TestMySQL(Validator):
self.validate_identity(
"INSERT INTO x VALUES (1, 'a', 2.0) ON DUPLICATE KEY UPDATE x.id = 1"
)
self.validate_identity(
"CREATE OR REPLACE VIEW my_view AS SELECT column1 AS `boo`, column2 AS `foo` FROM my_table WHERE column3 = 'some_value' UNION SELECT q.* FROM fruits_table, JSON_TABLE(Fruits, '$[*]' COLUMNS(id VARCHAR(255) PATH '$.$id', value VARCHAR(255) PATH '$.value')) AS q",
)
self.validate_all(
"CREATE TABLE z (a INT) ENGINE=InnoDB AUTO_INCREMENT=1 DEFAULT CHARACTER SET=utf8 COLLATE=utf8_bin COMMENT='x'",

View file

@ -234,21 +234,30 @@ MATCH_RECOGNIZE (
def test_json_table(self):
self.validate_identity(
"SELECT * FROM JSON_TABLE(foo FORMAT JSON, 'bla' ERROR ON ERROR NULL ON EMPTY COLUMNS (foo PATH 'bar'))"
"SELECT * FROM JSON_TABLE(foo FORMAT JSON, 'bla' ERROR ON ERROR NULL ON EMPTY COLUMNS(foo PATH 'bar'))"
)
self.validate_identity(
"SELECT * FROM JSON_TABLE(foo FORMAT JSON, 'bla' ERROR ON ERROR NULL ON EMPTY COLUMNS foo PATH 'bar')",
"SELECT * FROM JSON_TABLE(foo FORMAT JSON, 'bla' ERROR ON ERROR NULL ON EMPTY COLUMNS (foo PATH 'bar'))",
"SELECT * FROM JSON_TABLE(foo FORMAT JSON, 'bla' ERROR ON ERROR NULL ON EMPTY COLUMNS(foo PATH 'bar'))",
)
self.validate_identity(
"""SELECT
CASE WHEN DBMS_LOB.GETLENGTH(info) < 32000 THEN DBMS_LOB.SUBSTR(info) END AS info_txt,
info AS info_clob
FROM schemaname.tablename ar
INNER JOIN JSON_TABLE(:emps, '$[*]' COLUMNS (empno NUMBER PATH '$')) jt
INNER JOIN JSON_TABLE(:emps, '$[*]' COLUMNS(empno NUMBER PATH '$')) jt
ON ar.empno = jt.empno""",
pretty=True,
)
self.validate_identity(
"""SELECT
*
FROM JSON_TABLE(res, '$.info[*]' COLUMNS(
tempid NUMBER PATH '$.tempid',
NESTED PATH '$.calid[*]' COLUMNS(last_dt PATH '$.last_dt ')
)) src""",
pretty=True,
)
def test_connect_by(self):
start = "START WITH last_name = 'King'"

View file

@ -6,6 +6,18 @@ class TestRedshift(Validator):
dialect = "redshift"
def test_redshift(self):
self.validate_all(
"SELECT APPROXIMATE COUNT(DISTINCT y)",
read={
"spark": "SELECT APPROX_COUNT_DISTINCT(y)",
},
write={
"redshift": "SELECT APPROXIMATE COUNT(DISTINCT y)",
"spark": "SELECT APPROX_COUNT_DISTINCT(y)",
},
)
self.validate_identity("SELECT APPROXIMATE AS y")
self.validate_identity(
"SELECT 'a''b'",
"SELECT 'a\\'b'",

View file

@ -361,7 +361,18 @@ TBLPROPERTIES (
"SELECT CAST(123456 AS VARCHAR(3))",
write={
"": "SELECT TRY_CAST(123456 AS TEXT)",
"databricks": "SELECT TRY_CAST(123456 AS STRING)",
"spark": "SELECT CAST(123456 AS STRING)",
"spark2": "SELECT CAST(123456 AS STRING)",
},
)
self.validate_all(
"SELECT TRY_CAST('a' AS INT)",
write={
"": "SELECT TRY_CAST('a' AS INT)",
"databricks": "SELECT TRY_CAST('a' AS INT)",
"spark": "SELECT TRY_CAST('a' AS INT)",
"spark2": "SELECT CAST('a' AS INT)",
},
)
self.validate_all(

View file

@ -48,6 +48,14 @@ class TestTeradata(Validator):
self.validate_identity("HELP STATISTICS personnel.employee FROM my_qcd")
def test_create(self):
self.validate_identity(
"REPLACE VIEW view_b (COL1, COL2) AS LOCKING ROW FOR ACCESS SELECT COL1, COL2 FROM table_b",
"CREATE OR REPLACE VIEW view_b (COL1, COL2) AS LOCKING ROW FOR ACCESS SELECT COL1, COL2 FROM table_b",
)
self.validate_identity(
"REPLACE VIEW view_b (COL1, COL2) AS LOCKING ROW FOR ACCESS SELECT COL1, COL2 FROM table_b",
"CREATE OR REPLACE VIEW view_b (COL1, COL2) AS LOCKING ROW FOR ACCESS SELECT COL1, COL2 FROM table_b",
)
self.validate_identity("CREATE TABLE x (y INT) PRIMARY INDEX (y) PARTITION BY y INDEX (y)")
self.validate_identity("CREATE TABLE x (y INT) PARTITION BY y INDEX (y)")
self.validate_identity(

View file

@ -970,19 +970,19 @@ WHERE
self.validate_all(
"TRY_CONVERT(NVARCHAR, x, 121)",
write={
"spark": "CAST(DATE_FORMAT(x, 'yyyy-MM-dd HH:mm:ss.SSSSSS') AS VARCHAR(30))",
"spark": "TRY_CAST(DATE_FORMAT(x, 'yyyy-MM-dd HH:mm:ss.SSSSSS') AS VARCHAR(30))",
},
)
self.validate_all(
"TRY_CONVERT(INT, x)",
write={
"spark": "CAST(x AS INT)",
"spark": "TRY_CAST(x AS INT)",
},
)
self.validate_all(
"TRY_CAST(x AS INT)",
write={
"spark": "CAST(x AS INT)",
"spark": "TRY_CAST(x AS INT)",
},
)
self.validate_all(

View file

@ -15,6 +15,35 @@ SELECT
"q"."x" AS "x"
FROM UNNEST(ARRAY(1, 2)) AS "q"("x", "y");
# title: explode_outer
# dialect: spark
# execute: false
CREATE OR REPLACE TEMPORARY VIEW latest_boo AS
SELECT
TRIM(split(points, ':')[0]) as points_type,
TRIM(split(points, ':')[1]) as points_value
FROM (
SELECT
explode_outer(split(object_pointsText, ',')) as points
FROM (
SELECT
object_pointstext,
FROM boo
)
WHERE object_pointstext IS NOT NULL
);
CREATE OR REPLACE TEMPORARY VIEW `latest_boo` AS
SELECT
TRIM(SPLIT(`_q_1`.`points`, ':')[0]) AS `points_type`,
TRIM(SPLIT(`_q_1`.`points`, ':')[1]) AS `points_value`
FROM (
SELECT
EXPLODE_OUTER(SPLIT(`boo`.`object_pointstext`, ',')) AS `points`
FROM `boo` AS `boo`
WHERE
NOT `boo`.`object_pointstext` IS NULL
) AS `_q_1`;
# title: Union in CTE
WITH cte AS (
(

View file

@ -70,6 +70,15 @@ WITH cte AS (SELECT 1 AS x, 3 AS z) SELECT cte.a AS a, cte.z AS z FROM cte AS ct
WITH cte(x, y, z) AS (SELECT 1, 2, 3) SELECT a, z FROM (SELECT * FROM cte AS cte(b)) AS cte(a);
WITH cte AS (SELECT 1 AS x, 3 AS z) SELECT cte.a AS a, cte.z AS z FROM (SELECT cte.b AS a, cte.z AS z FROM cte AS cte(b)) AS cte;
WITH y AS (SELECT a FROM x) SELECT 1 FROM y;
WITH y AS (SELECT 1 AS _ FROM x AS x) SELECT 1 AS "1" FROM y;
WITH y AS (SELECT SUM(a) FROM x) SELECT 1 FROM y;
WITH y AS (SELECT MAX(1) AS _ FROM x AS x) SELECT 1 AS "1" FROM y;
WITH y AS (SELECT a FROM x GROUP BY a) SELECT 1 FROM y;
WITH y AS (SELECT 1 AS _ FROM x AS x GROUP BY x.a) SELECT 1 AS "1" FROM y;
--------------------------------------
-- Unknown Star Expansion
--------------------------------------

View file

@ -625,7 +625,7 @@ t0.x = t1.x AND t0.y < t1.y AND t0.y <= t1.y;
t0.x = t1.x AND t0.y < t1.y AND t0.y <= t1.y;
--------------------------------------
-- Coalesce
-- COALESCE
--------------------------------------
COALESCE(x);
x;
@ -669,18 +669,45 @@ a AND b AND (ROW() OVER () = 1 OR ROW() OVER () IS NULL);
CONCAT(x, y);
CONCAT(x, y);
CONCAT_WS(sep, x, y);
CONCAT_WS(sep, x, y);
CONCAT(x);
x;
CONCAT('a', 'b', 'c');
'abc';
CONCAT('a', NULL);
CONCAT('a', NULL);
CONCAT_WS('-', 'a', 'b', 'c');
'a-b-c';
CONCAT('a', x, y, 'b', 'c');
CONCAT('a', x, y, 'bc');
CONCAT_WS('-', 'a', x, y, 'b', 'c');
CONCAT_WS('-', 'a', x, y, 'b-c');
'a' || 'b';
'ab';
CONCAT_WS('-', 'a');
'a';
CONCAT_WS('-', x, y);
CONCAT_WS('-', x, y);
CONCAT_WS('', x, y);
CONCAT_WS('', x, y);
CONCAT_WS('-', x);
CONCAT_WS('-', x);
CONCAT_WS(sep, 'a', 'b');
CONCAT_WS(sep, 'a', 'b');
'a' || 'b' || x;
CONCAT('ab', x);
@ -837,3 +864,60 @@ x < CAST('2020-01-07' AS DATE);
x - INTERVAL '1' day = CAST(y AS DATE);
x - INTERVAL '1' day = CAST(y AS DATE);
--------------------------------------
-- Constant Propagation
--------------------------------------
x = 5 AND y = x;
x = 5 AND y = 5;
5 = x AND y = x;
y = 5 AND 5 = x;
x = 5 OR y = x;
x = 5 OR y = x;
(x = 5 AND y = x) OR y = 1;
(x = 5 AND y = 5) OR y = 1;
t.x = 5 AND y = x;
t.x = 5 AND y = x;
t.x = 'a' AND y = CONCAT_WS('-', t.x, 'b');
t.x = 'a' AND y = 'a-b';
x = 5 AND y = x AND y + 1 < 5;
FALSE;
x = 5 AND x = 6;
FALSE;
x = 5 AND (y = x OR z = 1);
x = 5 AND (y = x OR z = 1);
x = 5 AND x + 3 = 8;
x = 5;
x = 5 AND (SELECT x FROM t WHERE y = 1);
x = 5 AND (SELECT x FROM t WHERE y = 1);
x = 1 AND y > 0 AND (SELECT z = 5 FROM t WHERE y = 1);
x = 1 AND y > 0 AND (SELECT z = 5 FROM t WHERE y = 1);
x = 1 AND x = y AND (SELECT z FROM t WHERE a AND (b OR c));
x = 1 AND (SELECT z FROM t WHERE a AND (b OR c)) AND 1 = y;
t1.a = 39 AND t2.b = t1.a AND t3.c = t2.b;
t1.a = 39 AND t2.b = 39 AND t3.c = 39;
x = 1 AND CASE WHEN x = 5 THEN FALSE ELSE TRUE END;
x = 1 AND CASE WHEN FALSE THEN FALSE ELSE TRUE END;
x = 1 AND IF(x = 5, FALSE, TRUE);
x = 1 AND CASE WHEN FALSE THEN FALSE ELSE TRUE END;
x = y AND CASE WHEN x = 5 THEN FALSE ELSE TRUE END;
x = y AND CASE WHEN x = 5 THEN FALSE ELSE TRUE END;
x = 1 AND CASE WHEN y = 5 THEN x = z END;
x = 1 AND CASE WHEN y = 5 THEN 1 = z END;

View file

@ -2029,18 +2029,33 @@ JOIN "date_dim" AS "date_dim"
ON "date_dim"."d_year" = 2001
AND "store_sales"."ss_sold_date_sk" = "date_dim"."d_date_sk"
JOIN "household_demographics" AS "household_demographics"
ON "customer_demographics"."cd_demo_sk" = "store_sales"."ss_cdemo_sk"
ON (
"customer_demographics"."cd_demo_sk" = "store_sales"."ss_cdemo_sk"
AND "customer_demographics"."cd_education_status" = 'Advanced Degree'
AND "customer_demographics"."cd_education_status" = 'Primary'
AND "customer_demographics"."cd_education_status" = 'Secondary'
AND "customer_demographics"."cd_marital_status" = 'D'
AND "customer_demographics"."cd_marital_status" = 'M'
AND "customer_demographics"."cd_marital_status" = 'U'
AND "household_demographics"."hd_dep_count" = 1
AND "household_demographics"."hd_dep_count" = 3
AND "store_sales"."ss_hdemo_sk" = "household_demographics"."hd_demo_sk"
AND "store_sales"."ss_sales_price" <= 150.00
AND "store_sales"."ss_sales_price" >= 100.00
)
OR (
"customer_demographics"."cd_demo_sk" = "store_sales"."ss_cdemo_sk"
AND "customer_demographics"."cd_education_status" = 'Primary'
AND "customer_demographics"."cd_marital_status" = 'M'
AND "household_demographics"."hd_dep_count" = 1
AND "store_sales"."ss_hdemo_sk" = "household_demographics"."hd_demo_sk"
AND "store_sales"."ss_sales_price" <= 100.00
AND "store_sales"."ss_sales_price" >= 50.00
)
OR (
"customer_demographics"."cd_demo_sk" = "store_sales"."ss_cdemo_sk"
AND "customer_demographics"."cd_education_status" = 'Secondary'
AND "customer_demographics"."cd_marital_status" = 'D'
AND "household_demographics"."hd_dep_count" = 1
AND "store_sales"."ss_hdemo_sk" = "household_demographics"."hd_demo_sk"
AND "store_sales"."ss_sales_price" <= 200.00
AND "store_sales"."ss_sales_price" >= 150.00
)
JOIN "store" AS "store"
ON "store"."s_store_sk" = "store_sales"."ss_store_sk";

View file

@ -4,6 +4,7 @@ from datetime import date
from multiprocessing import Pool
import duckdb
import numpy as np
import pandas as pd
from pandas.testing import assert_frame_equal
@ -94,6 +95,11 @@ class TestExecutor(unittest.TestCase):
sql, _ = self.sqls[i]
a = self.cached_execute(sql)
b = pd.DataFrame(table.rows, columns=table.columns)
# The executor represents NULL values as None, whereas DuckDB represents them as NaN,
# and so the following is done to silence Pandas' "Mismatched null-like values" warnings
b = b.fillna(value=np.nan)
assert_frame_equal(a, b, check_dtype=False, check_index_type=False)
def test_execute_callable(self):

View file

@ -45,6 +45,10 @@ def normalize(expression, **kwargs):
return optimizer.simplify.simplify(expression)
def simplify(expression, **kwargs):
return optimizer.simplify.simplify(expression, constant_propagation=True, **kwargs)
class TestOptimizer(unittest.TestCase):
maxDiff = None
@ -271,7 +275,7 @@ class TestOptimizer(unittest.TestCase):
self.check_file("pushdown_projections", pushdown_projections, schema=self.schema)
def test_simplify(self):
self.check_file("simplify", optimizer.simplify.simplify)
self.check_file("simplify", simplify)
expression = parse_one("TRUE AND TRUE AND TRUE")
self.assertEqual(exp.true(), optimizer.simplify.simplify(expression))
@ -823,6 +827,11 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
self.assertEqual(exp.DataType.Type.ARRAY, expression.selects[0].type.this)
self.assertEqual(expression.selects[0].type.sql(), "ARRAY<INT>")
schema = MappingSchema({"t": {"c": "STRUCT<`f` STRING>"}}, dialect="bigquery")
expression = annotate_types(parse_one("SELECT t.c FROM t"), schema=schema)
self.assertEqual(expression.selects[0].type.sql(dialect="bigquery"), "STRUCT<`f` STRING>")
def test_type_annotation_cache(self):
sql = "SELECT 1 + 1"
expression = annotate_types(parse_one(sql))

View file

@ -272,3 +272,8 @@ class TestSchema(unittest.TestCase):
str(ctx.exception),
"Table z must match the schema's nesting level: 2.",
)
def test_has_column(self):
schema = MappingSchema({"x": {"c": "int"}})
self.assertTrue(schema.has_column("x", exp.column("c")))
self.assertFalse(schema.has_column("x", exp.column("k")))