diff --git a/README.md b/README.md index fccf9f7..66be572 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,8 @@ print(sql) The transpiled SQL can be executed with fast genome-unaware databases or in-memory analytic engines like DuckDB. +By default a column-to-column `INTERSECTS` join emits the naive overlap predicate (`a.chrom = b.chrom AND a.start < b.end AND b.start < a.end`) as a plain `ON` condition, which each engine's optimizer plans as a range join. For DuckDB you can additionally pass `dialect="duckdb"` to opt into a per-chromosome IEJoin plan for INNER, SEMI, or ANTI joins; shapes it declines fall through to the naive predicate. See [DuckDB IEJoin Dialect](https://giql.readthedocs.io/en/latest/transpilation/performance.html#duckdb-iejoin-dialect) for the supported shapes and fallback rules. + You can also use [oxbow](https://oxbow.readthedocs.io) to efficiently stream specialized genomics formats into DuckDB. ```python diff --git a/demo.ipynb b/demo.ipynb index 71729f3..5d5573e 100644 --- a/demo.ipynb +++ b/demo.ipynb @@ -224,7 +224,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [ { @@ -3579,6 +3579,152 @@ "result.head(15)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## 8. DISJOIN operator\n", + "\n", + "**Split intervals at breakpoints** so that no resulting sub-interval ever partially overlaps a reference interval.\n", + "\n", + "Unlike MERGE and CLUSTER, which *collapse* intervals, DISJOIN *multiplies* rows: one input interval becomes one or more sub-interval rows. It is a table-valued function, so it appears in the `FROM` clause like NEAREST. The full input row passes through unchanged and the sub-interval is appended as `disjoin_chrom`, `disjoin_start`, and `disjoin_end`.\n", + "\n", + "In **self-mode**, `DISJOIN(features)` splits a set against its own breakpoints. Selecting the distinct sub-intervals yields the globally non-overlapping partition — the equivalent of Bioconductor's `disjoin()`.\n", + "\n", + "**GIQL Query**: `SELECT DISTINCT disjoin_chrom, disjoin_start, disjoin_end FROM DISJOIN(features_a)`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "giql_query = \"\"\"\n", + " SELECT DISTINCT disjoin_chrom, disjoin_start, disjoin_end\n", + " FROM DISJOIN(features_a)\n", + " ORDER BY disjoin_chrom, disjoin_start\n", + "\"\"\"\n", + "\n", + "sql = transpile(giql_query, tables=[\"features_a\"])\n", + "print(\"Transpiled SQL:\")\n", + "print(sqlglot.transpile(sql, pretty=True)[0])\n", + "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n", + "\n", + "n_original = conn.execute(\"SELECT COUNT(*) FROM features_a\").fetchone()[0]\n", + "result = conn.execute(sql).fetchdf()\n", + "print(\n", + " f\"Result: {len(result)} non-overlapping segments \"\n", + " f\"(from {n_original} original peaks)\"\n", + ")\n", + "result.head(10)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Disjoining a set against itself\n", + "\n", + "The example above runs self-mode at scale; on a small, deliberately overlapping set every cut is visible. With no `reference`, `DISJOIN` defaults the reference to the target set, so each interval is split at every *other* interval's interior breakpoint: `A = [0, 100)` is cut where `B` starts (50); `B = [50, 150)` is cut at `A`'s end (100) and `C`'s start (120); `D = [0, 80)` is cut where `E` starts (40). Selecting `DISTINCT disjoin_*` collapses the pieces into the non-overlapping partition.\n", + "\n", + "**GIQL Query**: `SELECT * FROM DISJOIN(toy_peaks)`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# A small set of deliberately overlapping intervals, so the split is visible.\n", + "conn.execute(\"DROP TABLE IF EXISTS toy_peaks\")\n", + "conn.execute(\n", + " 'CREATE TABLE toy_peaks(chrom VARCHAR, \"start\" INTEGER, \"end\" INTEGER, name VARCHAR)'\n", + ")\n", + "conn.executemany(\n", + " \"INSERT INTO toy_peaks VALUES (?, ?, ?, ?)\",\n", + " [\n", + " (\"chr1\", 0, 100, \"A\"),\n", + " (\"chr1\", 50, 150, \"B\"),\n", + " (\"chr1\", 120, 200, \"C\"),\n", + " (\"chr2\", 0, 80, \"D\"),\n", + " (\"chr2\", 40, 80, \"E\"),\n", + " ],\n", + ")\n", + "print(\"Input intervals (overlapping):\")\n", + "print(\n", + " conn.execute(\n", + " 'SELECT name, chrom, start, \"end\" FROM toy_peaks ORDER BY chrom, start'\n", + " ).fetchdf()\n", + ")\n", + "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n", + "\n", + "# With no reference, DISJOIN splits toy_peaks against its own breakpoints.\n", + "giql_query = \"\"\"\n", + " SELECT name, disjoin_chrom, disjoin_start, disjoin_end\n", + " FROM DISJOIN(toy_peaks)\n", + " ORDER BY disjoin_chrom, disjoin_start, name\n", + "\"\"\"\n", + "\n", + "sql = transpile(giql_query, tables=[\"toy_peaks\"])\n", + "print(\"Transpiled SQL:\")\n", + "print(sqlglot.transpile(sql, pretty=True)[0])\n", + "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n", + "\n", + "result = conn.execute(sql).fetchdf()\n", + "print(\n", + " f\"Result: {len(result)} sub-intervals -- each interval split at its \"\n", + " f\"neighbours' breakpoints\"\n", + ")\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Splitting against a reference set\n", + "\n", + "Pass a second interval set as `reference :=` to cut each target feature at that set's breakpoints, keeping only the pieces a reference interval covers. Here every peak in `features_a` is split at the breakpoints of `features_b`; the parent peak coordinates ride along with each sub-interval, and a peak that overlaps no feature in B drops out entirely.\n", + "\n", + "**GIQL Query**: `SELECT * FROM DISJOIN(features_a, reference := features_b)`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "giql_query = \"\"\"\n", + " SELECT\n", + " chrom,\n", + " start AS peak_start,\n", + " \"end\" AS peak_end,\n", + " signal_value,\n", + " disjoin_start,\n", + " disjoin_end\n", + " FROM DISJOIN(features_a, reference := features_b)\n", + " ORDER BY chrom, disjoin_start\n", + "\"\"\"\n", + "\n", + "sql = transpile(giql_query, tables=[\"features_a\", \"features_b\"])\n", + "print(\"Transpiled SQL:\")\n", + "print(sqlglot.transpile(sql, pretty=True)[0])\n", + "print(\"\\n\" + \"=\" * 80 + \"\\n\")\n", + "\n", + "n_original = conn.execute(\"SELECT COUNT(*) FROM features_a\").fetchone()[0]\n", + "result = conn.execute(sql).fetchdf()\n", + "print(\n", + " f\"Result: {len(result)} sub-intervals of A split at B's breakpoints \"\n", + " f\"(from {n_original} original peaks)\"\n", + ")\n", + "result.head(10)" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -3587,7 +3733,7 @@ "\n", "## Summary\n", "\n", - "This demo showcased all 7 GIQL operators:\n", + "This demo showcased all 8 GIQL operators:\n", "\n", "1. **INTERSECTS**: Binary predicate for overlapping intervals\n", "2. **WITHIN**: Binary predicate for containment (A within B)\n", @@ -3596,6 +3742,7 @@ "5. **CLUSTER**: Aggregation operator to assign cluster IDs to overlapping intervals\n", "6. **DISTANCE**: UDF operator to calculate genomic distances between intervals\n", "7. **NEAREST**: Table-valued function for finding k-nearest genomic features\n", + "8. **DISJOIN**: Set operator to split intervals at reference breakpoints\n", "\n", "Each operator was:\n", "- Written in GIQL syntax\n", @@ -3608,7 +3755,7 @@ ], "metadata": { "kernelspec": { - "display_name": "giql", + "display_name": "giql (3.13.3)", "language": "python", "name": "python3" }, @@ -3622,7 +3769,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.2" + "version": "3.13.3" } }, "nbformat": 4, diff --git a/docs/dialect/aggregation-operators.rst b/docs/dialect/aggregation-operators.rst index 9887b87..ae57ef6 100644 --- a/docs/dialect/aggregation-operators.rst +++ b/docs/dialect/aggregation-operators.rst @@ -40,6 +40,9 @@ Syntax -- Strand-specific clustering CLUSTER(interval, stranded := true) AS cluster_id + -- Predicate-gated clustering (run-length encoding on a column) + CLUSTER(interval, predicate := depth = PREV(depth)) AS cluster_id + -- Combined parameters CLUSTER(interval, distance, stranded := true) AS cluster_id @@ -56,6 +59,36 @@ Parameters **stranded** *(optional)* When ``true``, only cluster intervals on the same strand. Default: ``false``. +**predicate** *(optional)* + A boolean expression evaluated between each interval and its sorted + predecessor. When supplied, the cluster-boundary condition becomes + **adjacent AND predicate**: an interval stays in the current cluster only + when it is within ``distance`` of its predecessor *and* the predicate holds + between the two. A change in the predicate forces a new cluster, so an + equality predicate yields a run-length encoding of the input sequence. + Omitting the predicate preserves the default adjacency-only behavior. + + Bare column references resolve to the *current* interval; the predecessor's + value of a column is referenced with ``PREV(column)`` + (e.g. ``depth = PREV(depth)``). The predicate composes with ``distance`` and + ``stranded`` and is evaluated under the operator's existing per-chromosome + (and per-strand) partition and start-position order. + + Two constraints apply: + + - **References existing columns only.** The predicate *gates* merging on + columns already present on the input rows; it does not synthesize a + statistic. Coverage depth, for example, must already be a column on the + rows (typically produced upstream by :ref:`DISJOIN ` and + aggregation). + - **Pairwise only, with single-linkage drift.** The predicate compares each + interval to its immediate sorted predecessor (everything ``LAG`` can + express). Whole-cluster conditions are out of scope. When the predicate is + not an equivalence relation (e.g. ``ABS(score - PREV(score)) < 5``), + consecutive pairs may each satisfy it while the cluster's extremes do not + — the same single-linkage behavior that ``distance``-based clustering + already exhibits. + Return Value ~~~~~~~~~~~~ @@ -101,6 +134,19 @@ Cluster intervals separately by strand: FROM features ORDER BY chrom, strand, start +**Predicate-Gated Clustering:** + +Cut adjacent intervals into clusters wherever a column's value changes +(run-length encoding). ``PREV(column)`` references the predecessor row's value: + +.. code-block:: sql + + SELECT + *, + CLUSTER(interval, predicate := depth = PREV(depth)) AS cluster_id + FROM features + ORDER BY chrom, start + **Analyze Cluster Statistics:** Count features per cluster: @@ -145,6 +191,10 @@ Find regions with multiple overlapping features: INNER JOIN cluster_sizes s ON c.cluster_id = s.cluster_id WHERE s.size >= 3 +.. note:: + + **Synthesized flag hidden under** ``SELECT *``. A ``SELECT *, CLUSTER(...)`` query materializes an internal ``__giql_is_new_cluster`` flag in a subquery, so the outer star is emitted as ``SELECT * EXCEPT (__giql_is_new_cluster)`` to keep that helper column out of the result (#184). ``* EXCEPT`` is a DataFusion-family extension: the generic and ``datafusion`` dialects emit it, while ``duckdb`` spells the exclusion ``EXCLUDE``. Transpile with ``dialect="duckdb"`` to execute on DuckDB — the portable generic ``* EXCEPT`` form is not DuckDB-runnable. A qualified ``SELECT t.*, CLUSTER(...)`` receives the same treatment: because ``CLUSTER`` runs over a single relation, the qualifier is dropped and the outer star is emitted as the same bare ``* EXCEPT (__giql_is_new_cluster)`` (#185). An explicitly-projected ``CLUSTER`` (no star) surfaces no helper column and needs no exclusion. A star may also be combined with additional explicit projection items — ``SELECT *, 1 AS extra, CLUSTER(...)`` — and each item surfaces exactly once. An *aliased* sibling (``expr AS name``) is materialized once in the inner subquery under a reserved ``__giql_sibling_N`` name, EXCEPTed from the outer star so it is not re-surfaced, and re-projected — aliased back to ``name`` — at its own position, so its output-column position is preserved even when the ``CLUSTER`` sits between the star and the item (``SELECT *, CLUSTER(...) AS cid, 1 AS extra`` keeps ``cid`` before ``extra``). Using a reserved inner name (rather than the user's alias) keeps the mechanism correct even when ``name`` collides with a base column the star surfaces (``SELECT *, score * 10 AS score`` yields both the base ``score`` and the computed one, matching the un-clustered projection) or with another sibling's alias (``SELECT *, 1 AS x, 2 AS x``). A *non-aliased* sibling (a bare column already covered by the star, or an unnamed expression) is materialized the same way, under its own reserved ``__giql_sibling_N`` name, and re-projected at its written slot — keeping its own column name (a bare column) or its rendered text (an unnamed expression) — so every sibling's output-column position matches the identical projection without ``CLUSTER`` (#190). + Performance Notes ~~~~~~~~~~~~~~~~~ @@ -194,6 +244,9 @@ Syntax -- Strand-specific merge SELECT MERGE(interval, stranded := true) FROM features + -- Predicate-gated merge (merge only equal-valued adjacent runs) + SELECT MERGE(interval, predicate := depth = PREV(depth)) FROM features + -- Merge with additional aggregations SELECT MERGE(interval), @@ -214,6 +267,16 @@ Parameters **stranded** *(optional)* When ``true``, merge intervals separately by strand. Default: ``false``. +**predicate** *(optional)* + A boolean expression that further restricts which adjacent intervals are + merged. ``MERGE`` decomposes into :ref:`CLUSTER ` plus a + ``GROUP BY`` over the cluster id, so it inherits predicate-aware boundaries + directly — see the :ref:`CLUSTER predicate ` description + for the full semantics, the ``PREV(column)`` convention, the + references-existing-columns-only constraint, and the pairwise-only / + single-linkage caveat. Omitting the predicate preserves the default + adjacency-only merge. + Return Value ~~~~~~~~~~~~ @@ -256,6 +319,24 @@ Merge intervals separately by strand: SELECT MERGE(interval, stranded := true) FROM features +**Predicate-Gated Merge (coverage depth):** + +Merge only adjacent intervals that share the same coverage depth, reconstructing +a re-clustered, depth-segmented partition from per-breakpoint segments produced +by :ref:`DISJOIN ` and aggregation: + +.. code-block:: sql + + SELECT MERGE(interval, predicate := depth = PREV(depth)) + FROM ( + SELECT disjoin_chrom AS chrom, + disjoin_start AS start, + disjoin_end AS end, + COUNT(*) AS depth + FROM DISJOIN(features) + GROUP BY disjoin_chrom, disjoin_start, disjoin_end + ) AS segments + **Merge with Feature Count:** Count how many features were merged into each region: @@ -324,6 +405,59 @@ Notes - MERGE is an aggregate operation that processes all matching rows - The operation sorts data internally, so pre-sorting is not required +.. note:: + + CLUSTER and MERGE cannot be combined in a single ``SELECT`` — MERGE + aggregates rows away while CLUSTER is a per-row window over those same rows, + so no single query expresses both. Transpiling ``SELECT MERGE(interval), + CLUSTER(interval) FROM features`` raises a ``ValueError``. Use them in + separate queries instead — for example, CLUSTER over a subquery, or MERGE + over one. + +.. note:: + + MERGE cannot be projected alongside a star. MERGE aggregates rows into one + row per merged region, so a ``SELECT *, MERGE(...)`` or ``SELECT t.*, + MERGE(...)`` has no coherent per-row meaning — the star names the + pre-aggregation input columns of a relation MERGE has already collapsed and + grouped away. Transpiling either shape raises a ``ValueError`` rather than + emitting non-executable SQL (a bare ``*`` re-surfaces non-grouped columns + under the synthesized ``GROUP BY``; a qualified ``rel.*`` dangles an alias the + aggregation no longer exposes). Drop the star, or project only grouping + columns and aggregates (e.g. ``COUNT(*)``) alongside ``MERGE`` — not raw input + columns, which are neither grouped nor aggregated. This is unlike + :ref:`CLUSTER `, a per-row window over which a star *is* + meaningful and supported. + +.. note:: + + **Projecting columns alongside** ``MERGE``. ``MERGE`` emits one row per merged + region and always projects the grouping keys (``chrom``, and ``strand`` when + ``stranded := true``) followed by ``MIN(start)`` / ``MAX(end)``. The other items in + a ``SELECT ..., MERGE(...)`` projection are reconciled against that grouping + (#192): an explicit grouping-key column (the ``chrom`` of the *Merge by Chromosome* + shape above) is emitted once, not duplicated; an aggregate (``COUNT(*)``, + ``AVG(score)``, ``STRING_AGG(name, ',')``) is computed per merged region; and an + expression over only grouping-key columns (``UPPER(chrom)``) is kept. A raw, + non-aggregated column that is neither a grouping key nor derived from one (e.g. + ``score``, or ``start`` within ``MAX(score) + start``) has no coherent value per + merged region and raises a ``ValueError`` — as does a *window* aggregate over such a + column (``SUM(score) OVER (...)``), since a window is evaluated after the grouping and + does not collapse it. An item whose output name collides with the ``chrom`` / + ``start`` / ``end`` columns ``MERGE`` synthesizes (e.g. ``chrom AS start``) also raises + rather than emitting two columns of that name — alias it to a distinct name instead. + The collision check folds case (``chrom AS Start`` raises too, since SQL binds the + unquoted alias onto ``start``), while an unaliased expression over a grouping key + (``CAST(chrom AS VARCHAR)``) is kept — its emitted column name is the expression, not a + bare grouping-key name. Likewise, a ``GROUP BY`` may name only the grouping-key columns + ``MERGE`` groups by, as a *bare column reference*: ``GROUP BY chrom`` is honored (it is + subsumed by the synthesized per-merge grouping), while ``GROUP BY`` over any other + column — or an expression (``GROUP BY UPPER(chrom)``) or ordinal (``GROUP BY 1``) that + ``MERGE`` cannot map to its grouping — raises a ``ValueError`` rather than being + silently discarded. A ``HAVING`` is honored against the per-merge grouping + (``HAVING COUNT(*) > 1`` keeps only regions built from more than one input interval) + rather than being dropped. + Related Operators ~~~~~~~~~~~~~~~~~ diff --git a/docs/dialect/distance-operators.rst b/docs/dialect/distance-operators.rst index 5486ca9..551b2fc 100644 --- a/docs/dialect/distance-operators.rst +++ b/docs/dialect/distance-operators.rst @@ -16,10 +16,11 @@ Description ~~~~~~~~~~~ The ``DISTANCE`` operator returns the number of base pairs separating two genomic -intervals. It follows standard genomic distance conventions: +intervals, matching ``bedtools closest -d`` semantics: - **Overlapping intervals**: Returns ``0`` -- **Non-overlapping intervals**: Returns the gap in base pairs (positive integer) +- **Book-ended (adjacent) intervals** (``A.end == B.start`` in half-open coordinates): Returns ``1`` +- **Non-overlapping intervals**: Returns the half-open gap plus one (a raw gap of ``N`` bases reports ``N + 1``) - **Different chromosomes**: Returns ``NULL`` Syntax @@ -42,7 +43,8 @@ Return Value ~~~~~~~~~~~~ - ``0`` for overlapping intervals -- Positive integer (gap in base pairs) for non-overlapping same-chromosome intervals +- ``1`` for book-ended (adjacent) intervals, matching ``bedtools closest -d`` +- Positive integer (half-open gap ``+ 1``) for non-overlapping same-chromosome intervals - ``NULL`` for intervals on different chromosomes Examples @@ -307,6 +309,15 @@ Find nearby same-strand features within distance constraints: WHERE nearest.distance BETWEEN -10000 AND 10000 ORDER BY peaks.name, ABS(nearest.distance) +Target support +~~~~~~~~~~~~~~ + +A correlated ``NEAREST`` (its reference is an outer-row column) runs on lateral-capable engines — DuckDB and the generic target — via a correlated ``LATERAL`` subquery, and on Apache DataFusion, which has no correlated-``LATERAL`` physical plan, via a decorrelated window-function rewrite. The two forms return the same result set, including under ``SELECT *`` / ``SELECT b.*`` (on DataFusion the internal helper columns are projected away — see the note below): the ``(start, end)`` tiebreaker orders rows tied at the k-th distance the same way on every engine, deterministically whenever ``(start, end)`` distinguishes the tied candidates. A trailing top-level ``ORDER BY`` is preserved as a top-level ordering except under the DataFusion star-projection wrapper, which sinks it into the wrapped subquery — there, and only there, rely on the *row set* rather than its order (an explicitly-projected query gets no wrapper on any target, so its ordering is preserved). The helper columns are hidden even in the composed cases — two correlated ``NEAREST`` fallbacks in one query, and a correlated ``NEAREST`` whose reserved columns are re-surfaced by an enclosing ``SELECT *`` *outside* its own SELECT (e.g. a wrapping ``CLUSTER``) — because the wrapper re-locates its target join by a reserved marker that survives the CLUSTER/MERGE rewrite and mints its reserved column names per fallback (#172). A standalone ``NEAREST`` with a literal reference is uncorrelated and uses the same ordered, limited subquery on every target. + +.. note:: + + ``SELECT *`` **/** ``SELECT b.*`` **over a correlated NEAREST on DataFusion.** The decorrelated window-function rewrite must expose its reference-key and rank columns (``__giql_x__rk_*``, ``__giql_x__rn``, minted per fallback) on the rewritten join for the equi-join to resolve. To keep them out of user output, the DataFusion path wraps the enclosing ``SELECT`` in ``SELECT * EXCEPT () FROM (...)`` (the ``EXCEPT`` list is the explicit reserved column names) when a ``SELECT *`` / ``SELECT b.*`` would surface them, so those projections return the same columns as the DuckDB LATERAL form (#160). The wrapper finds its target join by a reserved ``meta`` marker rather than a captured reference, so it still fires when a later ``CLUSTER`` / ``MERGE`` rewrite transplants the join into a rebuilt query; and because each fallback's reserved columns carry a per-fallback tag, two correlated ``NEAREST`` in one query each get their own ``* EXCEPT`` (#172). Explicitly-projected queries never surface the reserved columns and get no wrapper. + Notes ~~~~~ diff --git a/docs/dialect/index.rst b/docs/dialect/index.rst index 8d70e9d..dad1e3d 100644 --- a/docs/dialect/index.rst +++ b/docs/dialect/index.rst @@ -98,6 +98,24 @@ Combine and cluster genomic intervals. See :doc:`aggregation-operators` for detailed documentation. +Set Operations +-------------- + +Split genomic intervals against another set of intervals. + +.. list-table:: + :header-rows: 1 + :widths: 20 50 30 + + * - Operator + - Description + - Example + * - :ref:`DISJOIN ` + - Split intervals at reference breakpoints into sub-intervals + - ``SELECT * FROM DISJOIN(features, reference := mask)`` + +See :doc:`set-operators` for detailed documentation. + Set Quantifiers --------------- @@ -127,4 +145,5 @@ See :doc:`quantifiers` for detailed documentation. .. spatial-operators .. distance-operators .. aggregation-operators +.. set-operators .. quantifiers diff --git a/docs/dialect/set-operators.rst b/docs/dialect/set-operators.rst new file mode 100644 index 0000000..17787c6 --- /dev/null +++ b/docs/dialect/set-operators.rst @@ -0,0 +1,166 @@ +Set Operations +============== + +Set operations cut a set of genomic intervals against another set of intervals, +producing finer-grained sub-intervals. Unlike the aggregation operators, which +collapse intervals into summarized regions, a set operation *multiplies* rows -- +one input interval becomes one or more output rows. + +.. _disjoin-operator: + +DISJOIN +------- + +Split genomic intervals at reference breakpoints into sub-intervals that +never partially overlap a reference interval. + +Description +~~~~~~~~~~~ + +The ``DISJOIN`` operator is a table function: it takes a set of *target* +intervals and a *reference* set of intervals, and cuts each target interval +at every reference breakpoint (a reference ``start`` or ``end`` position) +that falls strictly inside it. Every resulting sub-interval is fully +contained by each reference interval it overlaps -- it can never *partially* +overlap one. + +Unlike :ref:`MERGE ` and :ref:`CLUSTER `, +which aggregate intervals, ``DISJOIN`` *multiplies* rows: one target interval +becomes one or more sub-interval rows. The full target row passes through +unchanged and the sub-interval is appended as ``disjoin_chrom``, +``disjoin_start``, ``disjoin_end``. + +When no ``reference`` is given it defaults to the target set, so +``DISJOIN(features)`` splits the set against its own breakpoints. Selecting +the distinct sub-intervals then yields the globally non-overlapping partition +-- the equivalent of Bioconductor's ``GenomicRanges::disjoin()``. ``DISJOIN`` +deliberately departs from that set-reducing shape: as a SQL table function it +*multiplies* rows and keeps each parent row intact, so it composes with +``JOIN`` and ``SELECT``. ``SELECT DISTINCT disjoin_chrom, disjoin_start, +disjoin_end`` recovers the canonical ``disjoin()`` partition exactly. + +This is useful for: + +- Partitioning a set of intervals into non-overlapping segments +- Re-tiling target features onto a data-defined or external grid +- Splitting features so downstream aggregates never double-count overlaps + +Syntax +~~~~~~ + +.. code-block:: sql + + -- Self-mode: split the set against its own breakpoints + SELECT * FROM DISJOIN(features) + + -- Split target features against an explicit reference set + SELECT * FROM DISJOIN(features, reference := mask) + + -- The reference may be a subquery + SELECT * FROM DISJOIN(features, reference := (SELECT * FROM mask)) + +Parameters +~~~~~~~~~~ + +**target** + The table of intervals to split. Must be a table registered with the + transpiler so its genomic columns can be resolved. + +**reference** *(optional)* + A registered table, a CTE defined in the same query, or a ``(SELECT ...)`` + subquery whose interval boundaries supply the breakpoints. Defaults to + ``target`` when omitted. A bare name that is neither a registered table nor + an in-query CTE is rejected. + +Return Value +~~~~~~~~~~~~ + +Every column of the matched target row, passed through unchanged, plus the +sub-interval: + +- ``disjoin_chrom`` - Chromosome of the sub-interval +- ``disjoin_start`` - Start of the sub-interval +- ``disjoin_end`` - End of the sub-interval + +A sub-interval that overlaps no reference interval is dropped (the coverage +filter). In self-mode every sub-interval is covered by its own parent, so +nothing is dropped. In reference mode a target interval that overlaps no +reference interval at all yields no output rows -- ``DISJOIN`` can drop a +target entirely, not merely trim it. + +.. note:: + + ``disjoin_chrom`` / ``disjoin_start`` / ``disjoin_end`` are reserved output + column names. If the target table already carries a column with one of + those names, the output will contain two columns of that name; DuckDB + silently renames the second, so ``SELECT disjoin_start`` then resolves to + the passed-through parent column rather than the computed sub-interval. + Rename the conflicting target column before the call. + +Examples +~~~~~~~~ + +**Partition a set of intervals:** + +Given two overlapping intervals ``A = [0, 20)`` and ``B = [10, 30)``, +``DISJOIN`` in self-mode cuts ``A`` at breakpoint ``10`` and ``B`` at +breakpoint ``20``: + +.. code-block:: sql + + SELECT DISTINCT disjoin_chrom, disjoin_start, disjoin_end + FROM DISJOIN(features) + ORDER BY disjoin_start + + -- Returns the partition: [0,10), [10,20), [20,30) + +**Split features against a mask:** + +.. code-block:: sql + + SELECT name, disjoin_start, disjoin_end + FROM DISJOIN(features, reference := mask) + +**Re-tile against a uniform grid:** + +.. code-block:: sql + + WITH bins AS ( + SELECT 'chr1' AS chrom, x AS start, x + 1000 AS "end" + FROM range(0, 250000000, 1000) AS t(x) + ) + SELECT * FROM DISJOIN(features, reference := bins) + +.. note:: + + ``range()`` is DuckDB-specific table-generating syntax; on other engines + build the bin grid with the equivalent construct. The grid must span every + chromosome present in ``features`` -- a feature on a chromosome the grid + omits has no covering bin and is dropped by the coverage filter. + +.. note:: + + ``DISJOIN`` is a table function and appears in the ``FROM`` clause, like + ``NEAREST``. On PostgreSQL the derived table must be given an alias + (``FROM DISJOIN(features) AS d``). + +.. note:: + + The appended ``disjoin_start`` / ``disjoin_end`` columns are emitted in the + target table's coordinate system -- the same convention as its + passed-through ``start`` / ``end`` columns, so every column of an output row + shares one convention. Cut positions are computed canonically inside the + operator; only their final representation follows the target table. + +.. note:: + + ``DISJOIN`` operates on coordinates only and is strand-blind: a reference + interval cuts a target interval regardless of either one's strand. To + disjoin per strand, filter the target and reference to a single strand + before the call. + +Related Operators +~~~~~~~~~~~~~~~~~ + +- :ref:`MERGE ` - Collapse overlapping intervals into one +- :ref:`CLUSTER ` - Assign cluster IDs without splitting diff --git a/docs/dialect/spatial-operators.rst b/docs/dialect/spatial-operators.rst index df96fb6..3e0a107 100644 --- a/docs/dialect/spatial-operators.rst +++ b/docs/dialect/spatial-operators.rst @@ -99,12 +99,12 @@ Find all variants, with gene information where available: FROM variants v LEFT JOIN genes g ON v.interval INTERSECTS g.interval -Binned Join Internals -~~~~~~~~~~~~~~~~~~~~~ +Join Internals +~~~~~~~~~~~~~~ -Column-to-column ``INTERSECTS`` joins use a binned equi-join strategy internally: each interval is assigned to one or more fixed-width bins, and the join is performed on ``(chrom, bin)`` pairs. Because an interval that spans a bin boundary belongs to more than one bin, the raw binned join can produce duplicate matches for the same pair of source rows. +A column-to-column ``INTERSECTS`` join expands to the **naive overlap predicate** — ``a.chrom = b.chrom AND a.start < b.end AND b.start < a.end`` — left in place as an ordinary ``ON`` (or ``WHERE``) condition. GIQL emits standard SQL and lets each engine's optimizer plan the range join: on DuckDB and DataFusion this becomes a hash join keyed on ``chrom`` with the two position inequalities applied as a residual join filter. Because the predicate is a plain condition on the original tables, inner and outer joins (``LEFT`` / ``RIGHT`` / ``FULL``) get correct semantics — including unmatched-row preservation — with no special restructuring, and every output row's multiplicity follows standard SQL bag semantics with no ``DISTINCT`` workaround. -GIQL eliminates these duplicates inside a **pairs CTE** that computes the set of distinct ``(left_key, right_key)`` pairs from the binned inner join, then joins the original tables back through this CTE. Because ``DISTINCT`` is applied only to the key pairs — not to the output query — all source rows are preserved, including genuinely duplicate records. Overlap counts, aggregations, and standard SQL bag semantics all work correctly without requiring any workaround. +DuckDB additionally offers an opt-in ``dialect="duckdb"`` per-chromosome IEJoin plan for INNER / SEMI / ANTI joins that routes through its ``IE_JOIN`` / ``PIECEWISE_MERGE_JOIN`` family; shapes it declines fall through to the naive predicate. See the :doc:`performance guide ` for details. Related Operators ~~~~~~~~~~~~~~~~~ diff --git a/docs/index.rst b/docs/index.rst index 91c98d9..78af604 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -29,6 +29,7 @@ See the :doc:`GIQL dialect ` docs. dialect/spatial-operators dialect/distance-operators dialect/aggregation-operators + dialect/set-operators dialect/quantifiers Transpilation @@ -47,6 +48,7 @@ See the :doc:`GIQL transpiler ` docs. transpilation/schema-mapping transpilation/execution transpilation/performance + transpilation/extending transpilation/api-reference @@ -62,6 +64,7 @@ See the following :doc:`recipes ` to learn how to use GIQL effect recipes/intersect recipes/distance recipes/clustering + recipes/disjoin recipes/advanced recipes/bedtools-migration diff --git a/docs/recipes/clustering.rst b/docs/recipes/clustering.rst index 8703021..f56c85e 100644 --- a/docs/recipes/clustering.rst +++ b/docs/recipes/clustering.rst @@ -347,6 +347,86 @@ Compare raw vs merged coverage: **Use case:** Quantify the redundancy in your feature set. +Predicate-Gated Clustering and Merging +-------------------------------------- + +Both ``CLUSTER`` and ``MERGE`` accept an optional ``predicate :=`` argument that +further restricts which adjacent intervals are coalesced: an interval stays in +the current cluster only when it is adjacent to its predecessor *and* the +predicate holds between the two. Bare columns resolve to the current interval; +the predecessor's value is referenced with ``PREV(column)``. The predicate +references columns already present on the rows — it gates merging, it does not +synthesize a statistic — and it compares each interval only to its immediate +sorted predecessor (so non-equivalence predicates exhibit single-linkage drift, +just like ``distance``-based clustering). + +Run-Length Encoding on a Column +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Merge only adjacent intervals that share the same value, cutting a new region +wherever the value changes: + +.. code-block:: sql + + SELECT MERGE(interval, predicate := depth = PREV(depth)) + FROM segments + +**Use case:** Collapse a per-base or per-segment signal into maximal runs of +constant value (e.g. equal coverage depth, same genotype, same annotation). + +Run-Length Encode with CLUSTER +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Assign a distinct cluster id to each maximal equal-valued run while keeping the +individual rows: + +.. code-block:: sql + + SELECT + *, + CLUSTER(interval, predicate := depth = PREV(depth)) AS run_id + FROM segments + ORDER BY chrom, start + +**Use case:** Label run boundaries for inspection before aggregating. + +Reconstruct disjoin() Coverage Segments +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +GIQL's :ref:`DISJOIN ` primitive splits overlapping intervals +at every breakpoint but deliberately does not re-cluster the resulting +sub-intervals. Pairing it with a predicate-gated ``MERGE`` closes that gap: cut +the input at every breakpoint, aggregate coverage depth per segment, then merge +back the contiguous runs of equal depth — reproducing the re-clustered, +depth-annotated output Bioconductor's ``disjoin()`` produces: + +.. code-block:: sql + + SELECT MERGE(interval, predicate := depth = PREV(depth)) + FROM ( + SELECT disjoin_chrom AS chrom, + disjoin_start AS start, + disjoin_end AS end, + COUNT(*) AS depth + FROM DISJOIN(features) + GROUP BY disjoin_chrom, disjoin_start, disjoin_end + ) AS segments + +**Use case:** Build a re-clustered coverage profile from overlapping intervals, +the expression-based generalization of ``disjoin()`` to any pairwise condition. + +Multi-Column Predicate +~~~~~~~~~~~~~~~~~~~~~~~ + +Gate merging on more than one column by combining comparisons with ``AND``: + +.. code-block:: sql + + SELECT MERGE(interval, predicate := strand = PREV(strand) AND name = PREV(name)) + FROM features + +**Use case:** Keep merged regions homogeneous across several attributes at once. + Advanced Patterns ----------------- diff --git a/docs/recipes/disjoin.rst b/docs/recipes/disjoin.rst new file mode 100644 index 0000000..810416b --- /dev/null +++ b/docs/recipes/disjoin.rst @@ -0,0 +1,118 @@ +Disjoining Intervals +==================== + +This section covers patterns for splitting intervals at breakpoints using +GIQL's ``DISJOIN`` operator -- partitioning a set into non-overlapping +segments and re-tiling features against a reference grid. + +Partition a Set of Intervals +---------------------------- + +Build a Disjoint Partition +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Split a set of intervals into the maximal set of non-overlapping +sub-intervals defined by its own breakpoints: + +.. code-block:: sql + + SELECT DISTINCT disjoin_chrom, disjoin_start, disjoin_end + FROM DISJOIN(features) + ORDER BY disjoin_chrom, disjoin_start + +**Use case:** Given ChIP-seq peak calls pooled from several samples, produce a +non-overlapping segment track -- the equivalent of Bioconductor's ``disjoin()`` +-- so downstream signal or count aggregates never double-count a base covered +by two overlapping sample peaks. + +Track Each Segment's Parent +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Keep the parent feature alongside each sub-interval: + +.. code-block:: sql + + SELECT name, disjoin_start, disjoin_end + FROM DISJOIN(features) + ORDER BY name, disjoin_start + +**Use case:** See how each original feature was fragmented and which segment +came from which parent. + +Split Against a Reference +------------------------- + +Split Features Against a Mask +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Cut target features at the boundaries of a reference (mask) set, keeping only +the pieces the mask covers: + +.. code-block:: sql + + SELECT name, disjoin_start, disjoin_end + FROM DISJOIN(features, reference := mask) + +**Use case:** Restrict ATAC-seq or gene features to a set of callable (or +otherwise interesting) mask regions, splitting them at the mask boundaries so +no reported piece straddles a mask edge. + +Re-tile Against a Uniform Grid +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Pass a generated set of fixed-width bins as the reference: + +.. code-block:: sql + + WITH bins AS ( + SELECT 'chr1' AS chrom, x AS start, x + 1000 AS "end" + FROM range(0, 250000000, 1000) AS t(x) + ) + SELECT * FROM DISJOIN(features, reference := bins) + +**Use case:** Break features onto a uniform coordinate grid -- for example to +build a fixed-width binned coverage matrix -- so each piece falls within a +single bin. + +.. note:: + + ``range()`` is DuckDB-specific syntax for generating the bin grid; other + engines need their own generator. The grid must also span every chromosome + present in ``features``, or features on an uncovered chromosome are dropped. + +The CTE / Subquery Reference Contract +------------------------------------- + +When the ``reference`` operand is an in-query CTE or a ``(SELECT ...)`` +subquery (rather than a registered table), it **must** project the canonical +``chrom`` / ``start`` / ``end`` columns -- the 0-based, half-open genomic +coordinates the operator references. A registered table may declare a custom +column mapping or coordinate system, but a CTE or subquery carries no such +declaration, so GIQL assumes it is already canonical. + +GIQL validates this contract at transpile time. If a CTE or subquery operand +omits one of the canonical columns, ``transpile()`` raises a ``ValueError`` +naming the operator, the offending relation, and the missing column(s) -- a +clear GIQL-level diagnostic rather than an opaque engine ``column not found`` +error at execution time. For example, ``reference := (SELECT chrom, start FROM +mask)`` is rejected because it omits ``end``. + +.. code-block:: sql + + -- Valid: the subquery projects all three canonical columns + SELECT * FROM DISJOIN(features, reference := ( + SELECT chrom, start, "end" FROM mask WHERE score > 10 + )) + +This check is best-effort. When the operand's output columns cannot be +determined statically -- a ``SELECT *`` star projection, or a placeholder such +as ``SELECT 1`` -- the operand is passed through without a diagnostic (the +relation's schema is unknown at transpile time), and a genuinely missing column +would still surface as an engine error at execution time. + +Coming from Bedtools? +--------------------- + +``DISJOIN`` has no direct bedtools equivalent -- no single bedtools command +splits intervals at breakpoints. See the :doc:`bedtools-migration` guide for +the GIQL operators that do map onto bedtools commands. diff --git a/docs/recipes/index.rst b/docs/recipes/index.rst index cc97e47..173ad53 100644 --- a/docs/recipes/index.rst +++ b/docs/recipes/index.rst @@ -19,6 +19,10 @@ Recipe Categories Clustering overlapping intervals, distance-based clustering, merging intervals, and aggregating cluster statistics. +:doc:`disjoin` + Splitting intervals at breakpoints, partitioning a set into + non-overlapping segments, and re-tiling features against a reference grid. + :doc:`advanced` Multi-range matching, complex filtering with joins, aggregate statistics, window expansions, and multi-table queries. diff --git a/docs/transpilation/api-reference.rst b/docs/transpilation/api-reference.rst index fcba984..ca771a3 100644 --- a/docs/transpilation/api-reference.rst +++ b/docs/transpilation/api-reference.rst @@ -1,13 +1,48 @@ API Reference ============= -.. currentmodule:: giql +Transpilation +------------- -.. autosummary:: +.. autofunction:: giql.transpile.transpile - transpile - Table +.. autoclass:: giql.table.Table -.. autofunction:: transpile +Extension hook +-------------- -.. autoclass:: Table +The registry-based extension point for custom targets and operator expanders +(see :doc:`extending`). Every symbol below is re-exported from the top-level +``giql`` package. + +.. autofunction:: giql.expander.register + +.. autodata:: giql.expander.REGISTRY + +.. autoclass:: giql.expander.ExpanderRegistry + :members: register, register_target, target, resolve, has_override, + unregister, clear, snapshot, restore + +.. autoclass:: giql.expander.RegistrySnapshot + :members: + +.. autoclass:: giql.expander.ExpansionContext + :members: + +.. autoclass:: giql.expander.OperatorExpander + :members: + +Targets and capabilities +------------------------ + +.. autoclass:: giql.targets.Target + :members: + +.. autoclass:: giql.targets.Capabilities + :members: + +.. autoclass:: giql.targets.GenericTarget + +.. autoclass:: giql.targets.DuckDBTarget + +.. autoclass:: giql.targets.DataFusionTarget diff --git a/docs/transpilation/extending.rst b/docs/transpilation/extending.rst new file mode 100644 index 0000000..184cd28 --- /dev/null +++ b/docs/transpilation/extending.rst @@ -0,0 +1,201 @@ +Extending GIQL: custom targets and operators +============================================= + +Every GIQL operator is expanded to standard SQL by a small **expander** function +selected from a process-wide registry, keyed by ``(target, operator)``. That +registry is a supported public extension point. Through it you can: + +- **register a custom target** — a new SQL engine with its own capabilities and + sqlglot output dialect — and select it with ``transpile(dialect="")``; and +- **override an operator's expander** for a target, so that operator emits SQL + tailored to that engine. + +Both hooks live in the ``giql.expander`` and ``giql.targets`` modules, and the key +symbols are re-exported from the top-level package: + +.. code-block:: python + + from giql import ( + transpile, + register, # decorator: register/override an operator expander + REGISTRY, # the process-wide plugin registry + Target, # base class for a custom target + Capabilities, # a target's feature set + GenericTarget, # the portable fallback target + ExpansionContext, + ) + + +Targets and capabilities +------------------------ + +A :class:`~giql.targets.Target` is a first-class SQL engine. It carries a +``name``, the ``sqlglot_dialect`` used to serialize the final AST (``None`` +selects sqlglot's portable default), and a :class:`~giql.targets.Capabilities` +set. Capabilities — not scattered ``if dialect == ...`` branches — drive the +portable emission choices: + +- ``supports_lateral`` — correlated ``LATERAL`` joins (NEAREST uses a decorrelated + window-function form when this is ``False``); +- ``supports_star_replace`` — ``SELECT * REPLACE (...)`` (used for coordinate + canonicalization and the DISJOIN / NEAREST passthroughs; the portable + ``SELECT * EXCEPT (...)`` form is emitted otherwise); +- ``supports_qualify`` — the ``QUALIFY`` clause. + +The column-to-column INTERSECTS *join* strategy is **not** a capability flag: it is +selected by the operator-expander registry. Every target emits the naive overlap +predicate (the ``(GenericTarget, Intersects)`` expander) unless it registers a +target-specific ``(YourTarget, Intersects)`` override — as DuckDB does for its +per-chromosome IEJoin plan (:mod:`giql.expanders.intersects_duckdb`). + +Define a custom target by subclassing :class:`~giql.targets.Target` as a frozen +dataclass. Give every field a default so the class is constructible with no +arguments (the ``@register`` decorator instantiates a target class for you): + +.. code-block:: python + + from dataclasses import dataclass + from giql import Target, Capabilities + + @dataclass(frozen=True) + class PostgresTarget(Target): + name: str = "postgres" + sqlglot_dialect: str | None = "postgres" + capabilities: Capabilities = Capabilities( + supports_lateral=True, + supports_star_replace=False, # -> portable "* EXCEPT" canonicalization + supports_qualify=True, + ) + + +Registering and selecting a custom target +----------------------------------------- + +Declare the target on the registry, then select it by name. A **capability-only** +target — one that reuses every built-in operator expander and differs only in its +capabilities and dialect — needs nothing more than +:meth:`~giql.expander.ExpanderRegistry.register_target`: + +.. code-block:: python + + from giql import REGISTRY, transpile + + REGISTRY.register_target(PostgresTarget()) + + sql = transpile( + "SELECT * FROM peaks WHERE interval WITHIN 'chr1:1000-5000'", + tables=["peaks"], + dialect="postgres", + ) + +``dialect="postgres"`` resolves the registered target; its +``supports_star_replace=False`` capability selects the portable projection form, +and its ``sqlglot_dialect="postgres"`` serializes the final SQL. (The built-in +names ``"duckdb"`` and ``"datafusion"`` and ``None`` for the generic target still +resolve as before; ``"generic"`` is *not* a selectable name — ``None`` is the +sole way to select the generic target.) + + +Overriding an operator's expander +--------------------------------- + +To tailor one operator to an engine, register an expander for +``(target, operator)`` with the :func:`~giql.expander.register` decorator. An +expander takes the operator node and its :class:`~giql.expander.ExpansionContext` +and returns the sqlglot AST that replaces the node: + +.. code-block:: python + + from sqlglot import parse_one + from giql import register, ExpansionContext + from giql.expressions import Within + + @register(PostgresTarget, Within) + def within_between(node, ctx: ExpansionContext): + # ctx.resolution exposes the operator's resolved column operands, and + # ctx.capabilities / ctx.target the active engine's capabilities. + col = ctx.resolution.column("this") + lo, hi = 1000, 5000 # (a real expander parses the node's range) + return parse_one(f"{col.start} BETWEEN {lo} AND {hi}") + +Registering an expander also **declares its target by name** as a side effect, so +``@register(PostgresTarget, Within)`` alone makes ``dialect="postgres"`` +resolvable — a separate ``register_target`` call is only needed for a target that +overrides no operators. + +Resolution follows a fallback chain: an exact ``(target, operator)`` expander wins, +otherwise the built-in ``(GenericTarget(), operator)`` expander runs. So you only +register the operators that genuinely differ for your engine; everything else +reuses the portable built-ins. + +.. note:: + + ``ctx.resolution.column("this").start`` (and ``.end`` / ``.chrom``) return + **quoted** physical column names (e.g. ``"start"``, or ``a."start"`` when the + operand is aliased), ready to splice into a SQL string that you ``parse_one``. + Do not pass them to ``sqlglot.exp.column()``, + which would quote them a second time. Building fragments as strings and parsing + once — as the built-in expanders do — is the simplest correct approach. + + +The node-local boundary +----------------------- + +An expander's **return value** is node-local: ``expand(node, ctx)`` returns the +one expression that replaces the operator node in place. It cannot *return* a +reshaped enclosing query. An expander may still restructure the query it sits in +as a side effect and then return the node unchanged — the built-in CLUSTER and +MERGE expanders do exactly this, rewriting their single-table ``SELECT`` in place. + +When an expander must rewrite the **enclosing statement** — wrap an enclosing +``SELECT``, or reshape a projection it does not own — it registers a *statement +finalizer* via :meth:`~giql.expander.ExpansionContext.add_statement_finalizer`. +The pass applies every registered finalizer to the statement, in registration +order, after all node-local replacements complete; each receives the current +statement root and returns the (possibly new) root. The built-in NEAREST +DataFusion fallback uses this to wrap its output in +``SELECT * EXCEPT (...)`` and hide the reserved rank/key columns its decorrelated +join must expose: + +.. code-block:: python + + def expand(self, node, ctx): + # ... rewrite the node / enclosing join in place ... + ctx.add_statement_finalizer(lambda root: wrap_or_return(root)) + return node + +A finalizer's returned root is emitted **as-is** — the pass does not re-validate +it — so a finalizer that reshapes a projection must not reference columns or +relations absent from what it rewrites. Wrapping a projection in +``SELECT * EXCEPT (missing_col)``, for instance, builds without error at transpile +time but fails at engine runtime. The built-in fallback guards this by wrapping +only when the projection genuinely surfaces the columns it excepts; a custom +finalizer should apply the same discipline. + +A query-level fold that **adds or reshapes joins** across relations uses the same +finalizer seam. The DuckDB IEJoin plan for column-to-column INTERSECTS joins is +exactly this: its ``(DuckDBTarget, Intersects)`` override +(:mod:`giql.expanders.intersects_duckdb`) builds a whole-query per-chromosome +rewrite and registers a finalizer that replaces the statement root with it, rather +than replacing a single node in place (#169). + + +Undoing a registration +----------------------- + +The registry exposes a teardown seam so a plugin or a test fixture can undo its +registrations rather than reaching into private state: + +.. code-block:: python + + REGISTRY.unregister(PostgresTarget(), Within) # drop one expander entry + REGISTRY.clear() # drop all entries and targets + + saved = REGISTRY.snapshot() # save/restore around a body + try: + ... # register custom entries + finally: + REGISTRY.restore(saved) + +See :func:`giql.transpile.transpile`, :func:`giql.expander.register`, and +:class:`giql.targets.Capabilities` in the :doc:`api-reference`. diff --git a/docs/transpilation/index.rst b/docs/transpilation/index.rst index 6f4fcbc..7977be8 100644 --- a/docs/transpilation/index.rst +++ b/docs/transpilation/index.rst @@ -207,7 +207,7 @@ conventions yield the same numeric distance. FROM ( SELECT *, - SUM(is_new_cluster) OVER ( + SUM(__giql_is_new_cluster) OVER ( PARTITION BY "chrom" ORDER BY "start" NULLS LAST ) AS __giql_cluster_id @@ -220,9 +220,9 @@ conventions yield the same numeric distance. ORDER BY "start" NULLS LAST ) >= "start" THEN 0 ELSE 1 - END AS is_new_cluster + END AS __giql_is_new_cluster FROM features - ) AS lag_calc - ) AS clustered - GROUP BY chrom, __giql_cluster_id + ) AS __giql_lag_calc + ) AS __giql_clustered + GROUP BY "chrom", __giql_cluster_id ORDER BY "chrom" NULLS LAST, "start" NULLS LAST diff --git a/docs/transpilation/performance.rst b/docs/transpilation/performance.rst index 3815cd5..1fcd643 100644 --- a/docs/transpilation/performance.rst +++ b/docs/transpilation/performance.rst @@ -301,6 +301,303 @@ Analyze query execution plans by running EXPLAIN on the transpiled SQL: Backend-Specific Tips --------------------- +DuckDB IEJoin Dialect +~~~~~~~~~~~~~~~~~~~~~ + +By default (``dialect=None`` / ``"datafusion"``) a column-to-column +``INTERSECTS`` join emits the **naive overlap predicate** — a plain ``ON +a.chrom = b.chrom AND a.start < b.end AND b.start < a.end`` condition — and +lets the engine's own optimizer plan it. On DuckDB and DataFusion this +becomes a hash join keyed on ``chrom`` with the two position inequalities +as a residual join filter, correct for both inner and outer joins with no +special handling. This is the lowest-common-denominator plan: standard SQL +that every target runs. + +For column-to-column ``INTERSECTS`` joins (INNER, SEMI, or ANTI) on +DuckDB, the ``dialect="duckdb"`` opt-in instead transpiles the join into a +per-chromosome dynamic-SQL pattern. Each per-chromosome subquery contains +only inequality predicates, which DuckDB plans through its range-join +family (``IE_JOIN`` or ``PIECEWISE_MERGE_JOIN``). Shapes this path declines +fall through to the naive predicate above. + +.. code-block:: python + + from giql import transpile + + sql = transpile( + """ + SELECT a.chrom, a.start, b.start + FROM peaks a + JOIN genes b ON a.interval INTERSECTS b.interval + """, + tables=["peaks", "genes"], + dialect="duckdb", + ) + # sql is a two-statement script — see the next example for execution. + +The output is a multi-statement string of the form:: + + SET VARIABLE __giql_iejoin_ = COALESCE((... string_agg per chromosome ...), ''); + SELECT ... FROM query(getvariable('__giql_iejoin_')) AS __giql_iejoin_wrapper + +The ```` is a per-call ``uuid4().hex`` (128 bits) suffix so the +``SET VARIABLE`` name is collision-resistant even when outputs from many +``transpile()`` calls are interleaved in a single DuckDB session +(session variables are global session state). The outer SELECT's +wrapper-relation alias is constant (``__giql_iejoin_wrapper``) because +it isn't user-visible and doesn't need to vary per call. + +Inside the per-chromosome string-builder, chromosome names are emitted +via ``replace(chrom, '''', '''''')`` wrapped in single quotes, so +chromosome identifiers containing literal single quotes interpolate +safely into the dynamic SQL. + +Because the output is a two-statement script that depends on session +state, execute it through a single DuckDB connection's ``.execute()`` — +DuckDB ≥1.4 (the version pinned by GIQL's development and CI +environment) accepts multi-statement strings in one call and returns +the result of the final statement. SQLAlchemy's ``text()``, +``pandas.read_sql_query``, and similar wrappers that split or rewrite +the string may drop the ``SET VARIABLE`` and produce empty or NULL +results. + +.. code-block:: python + + import duckdb + from giql import transpile + + conn = duckdb.connect() + conn.execute("CREATE TABLE peaks (chrom VARCHAR, \"start\" INTEGER, \"end\" INTEGER)") + conn.execute("CREATE TABLE genes (chrom VARCHAR, \"start\" INTEGER, \"end\" INTEGER)") + sql = transpile( + """ + SELECT a.chrom, a.start, b.start + FROM peaks a + JOIN genes b ON a.interval INTERSECTS b.interval + """, + tables=["peaks", "genes"], + dialect="duckdb", + ) + rows = conn.execute(sql).fetchall() + +The dialect rewrites the whole query, so the supported shape is +``SELECT FROM +{INNER|SEMI|ANTI} JOIN {ON|WHERE} `` — the +JOIN and the column-to-column ``INTERSECTS`` are both required +(literal-range ``INTERSECTS`` against a single table falls through to +the default predicate generator). + +**Join variants.** ``INNER JOIN`` is the default. ``SEMI JOIN`` returns +distinct left rows that have at least one overlapping match; ``ANTI +JOIN`` returns left rows with no overlapping match. Both restrict the +outer SELECT to left-side projections — any ``b.col`` reference (or an +aggregate over ``b.col``) raises ``ValueError``; a right-side ``b.*`` +declines to the naive-predicate plan with every other star (which then +rejects the out-of-scope right table at bind time). ANTI uses a +left-distinct chromosome partition rather than the chromosome INTERSECT +used by INNER / SEMI, so left rows on chromosomes absent from the right +table are preserved. + +DuckDB plans all three variants through its ``IE_JOIN`` / +``PIECEWISE_MERGE_JOIN`` sort-merge range-join family. INNER emits a +per-chromosome inner join directly. SEMI and ANTI are emitted as a +correlated ``WHERE EXISTS`` / ``WHERE NOT EXISTS`` subquery over the +per-chromosome right partition, because DuckDB plans a bare ``SEMI +JOIN`` / ``ANTI JOIN`` inequality overlap as a quadratic +``BLOCKWISE_NL_JOIN`` (a nested loop) rather than the range-join fast +path — the correlated form reaches the fast path while preserving exact +semantics (#208). Expect large speedups over the naive predicate for all +three variants at scale. + +**count_overlaps.** The idiomatic per-interval overlap count — ``SELECT +a.cols, COUNT(b.col) FROM a LEFT JOIN b ON a.interval INTERSECTS +b.interval GROUP BY a.cols`` — is also accelerated (#209). A plain +``LEFT JOIN`` inequality overlap would decline to the naive predicate (a +``HASH_JOIN`` on ``chrom`` with the position inequalities as a residual +filter, quadratic when the chromosome key has low cardinality). Instead +the dialect computes the counts through the INNER ``IE_JOIN`` path and +wraps them in a zero-fill ``LEFT JOIN`` back onto the distinct left keys, +so left intervals with no overlap report ``0``. This covers a single +``COUNT()`` (or ``COUNT(DISTINCT )``) with a +``GROUP BY`` on the projected left columns; ``COUNT(*)``, non-COUNT +aggregates, ``ORDER BY`` / ``LIMIT`` / ``HAVING``, and other outer-join +shapes still take the naive-predicate plan. + +On top of the core shape the dialect also absorbs several common +decorations: + +- **Outer modifiers.** ``ORDER BY`` / ``LIMIT`` / ``OFFSET`` / + ``DISTINCT`` on the outer query are appended to the outer SELECT; + column references in ``ORDER BY`` are rewritten to the wrapper + relation's inner aliases. +- **Aggregates.** ``GROUP BY`` / ``HAVING`` with aggregate functions + (``COUNT``, ``SUM``, ``MIN``, ``MAX``, ``AVG``, ``COUNT(DISTINCT + a.col)``, ...) are appended to the outer SELECT. Aggregate arguments + must be table-qualified (``COUNT(*)`` and ``COUNT(a.col)`` are both + accepted). +- **Extra JOIN / WHERE predicates.** Additional non-INTERSECTS + predicates ANDed onto the join ON or WHERE (e.g. ``a.score > + b.score`` or ``WHERE a.score > 100``) are inlined into each + per-chromosome subquery's ON, so DuckDB filters them inside each + IEJoin candidate set. **Limitations:** the dialect peels extra + predicates only across top-level ``AND`` connectives, so predicates + that wrap the ``INTERSECTS`` in ``OR`` / ``NOT`` / parentheses still + fall back. Predicates whose AND-tree residual contains a subquery, + aggregate, or window function fall back too (DuckDB forbids window + functions inside ``JOIN ON``, and subqueries / aggregates inside an + IEJoin candidate set would either break the planner or produce + semantically wrong results). If you hit the + WHERE-INTERSECTS-plus-extra-JOIN-ON-predicate shape described in + `#94 `_ on the default + ``dialect=None`` path, ``dialect="duckdb"`` is one workaround — the + dialect inlines extras directly into each per-chromosome subquery + and is unaffected by that bug. + +The dialect splits unsupported shapes into two buckets. Soft-fallback +shapes route to the naive-predicate plan automatically and return correct +results; within those shapes the ``dialect`` kwarg is safe to set +without risk of silent incorrectness. Hard-error shapes (enumerated +further below) raise ``ValueError`` at transpile time — the dialect +deliberately refuses them rather than silently producing the wrong +SQL. + +The soft-fallback shapes are: + +- **Outer joins.** ``LEFT`` / ``RIGHT`` / ``FULL`` ``JOIN ... ON ... + INTERSECTS ...`` falls back to the naive-predicate plan. The same applies + when the join keeps its side modifier and the ``INTERSECTS`` lives in + the top-level ``WHERE`` (e.g. ``LEFT JOIN ... ON TRUE WHERE + a.interval INTERSECTS b.interval``). +- **SEMI / ANTI join with the INTERSECTS in the WHERE.** A ``SEMI`` / + ``ANTI`` join whose column-to-column ``INTERSECTS`` sits in the + top-level ``WHERE`` rather than its own ``ON`` (e.g. ``ANTI JOIN ... ON + TRUE WHERE a.interval INTERSECTS b.interval``) falls back. The right + table is out of scope in the ``WHERE`` after a left-only join, so the + reference plans reject it with a binder error; the dialect declines so + it surfaces that same error instead of relocating the predicate into + the join and inventing anti/semi-overlap results. The idiomatic form + with the ``INTERSECTS`` in the join ``ON`` is unaffected. +- **Star projections.** Any star in the top-level SELECT list — bare + ``*``, ``a.*``, ``b.*``, or ``a.* AS x`` — falls back. Building the + dialect's per-chromosome dynamic SQL requires enumerating the star's + columns at transpile time, but only the registered genomic columns + (``chrom`` / ``start`` / ``end`` / ``strand``) are known — arbitrary + user columns are invisible. The naive-predicate plan expands the real + star against DuckDB's live schema, so declining keeps the projection + identical across backends rather than silently narrowing it. +- **Projections the IEJoin cannot rebuild.** The dialect's projection + rebuild handles a qualified column (``a.col``) or a plain aggregate + over qualified columns (``SUM(a.score)`` / ``COUNT(*)``). Any other + SELECT-list shape the naive plan compiles for free — an expression + (``a.start + 1``), a window aggregate (``SUM(a.score) OVER (...)``), a + ``FILTER`` clause, a scalar subquery, an aggregate nested in an + expression (``COUNT(*) * 2``), a bare literal, or a star nested in an + aggregate argument (``COUNT(a.*)`` / ``MIN(COLUMNS(*))``) — falls back, + so ``dialect="duckdb"`` stays consistent with every other backend + instead of hard-erroring or, for the aggregate-nested star, + miscompiling. A projection whose column the rebuild cannot attribute to + a join side (unqualified, an unknown table, or the right side under a + SEMI / ANTI join) raises at transpile time instead (see the hard-error + list below) — the unknown-table and right-side cases are rejected by the + naive plan too, and a bare unqualified column, though it may be + naive-valid when unambiguous, has no side the dialect can infer without a + live schema. +- **NATURAL joins.** ``NATURAL JOIN`` falls back because the dialect + cannot enumerate shared columns at transpile time (only the + registered ``chrom`` / ``start`` / ``end`` / ``strand`` columns are + known). +- **USING joins.** Single-column ``USING()`` admits (the + per-chromosome partition is exactly the equi-join). Multi-column + ``USING`` and ``USING()`` fall back; inline support + is a documented follow-up. +- **Multiple INTERSECTS predicates.** Queries with more than one + column-to-column ``INTERSECTS`` (anywhere in the AST, including + subqueries) fall back. +- **INTERSECTS references unrecognized join sides.** A column-to-column + INTERSECTS whose operands don't reference the FROM table's alias, + or where the second operand's alias doesn't resolve to a registered + JOIN target, falls back to the naive-predicate plan. +- **Self-joins.** Joining a table to itself falls back. +- **OR / NOT / paren-wrapped extra predicates.** See the Extra JOIN / + WHERE predicates note above. +- **Subquery, aggregate, or window-function extra predicates.** + Predicates whose AND-tree residual contains a subquery, an aggregate + function, or a window function (``ROW_NUMBER() OVER (...)`` etc.) + fall back rather than inline into the per-chromosome subquery's + ``ON`` (DuckDB forbids window functions in ``ON`` clauses, and + aggregates / subqueries inside an IEJoin candidate set would either + break the planner or produce semantically wrong results). +- **Subquery inside GROUP BY / HAVING / ORDER BY.** A modifier clause + containing a nested subquery (including ``EXISTS (SELECT ...)``) + falls back to the naive-predicate plan, because the dialect's modifier-ref + rewriter is not scope-aware and would corrupt the subquery's + column scope. Subqueries hidden inside aggregate arguments in the + SELECT list fall back for the same reason. +- **Top-level WITH clauses.** A query with a leading ``WITH`` falls + back so the CTE survives. +- **More than two tables.** Any extra FROM-clause table (auxiliary + comma-join or non-INTERSECTS ``JOIN``) falls back. +- **Non-base-table operands.** A subquery, CTE, or GIQL table-function + (e.g. ``DISJOIN(genes)``) used as a join operand falls back. +- **DISTINCT ON (...).** Only plain ``DISTINCT`` is absorbed; the + ``DISTINCT ON`` variant falls back. + +Query-shape ``ValueError`` raises (the dialect deliberately refuses +these rather than silently miscompile): + +- **Unqualified columns.** A bare column reference (``SELECT score``) + raises — the dialect path needs to know which side each output column + comes from, and has no live schema to infer it (so an unambiguous + unqualified column that the naive plan would accept still raises here). + Use ``a.col`` or ``a.col AS x``. (Bare ``SELECT *`` does not raise; like + every star it falls back to the naive-predicate plan — see the + soft-fallback list above. An expression / window aggregate / ``FILTER`` + clause over an *unqualified* column raises the same "requires qualified + projections" error — once qualified, the wrapper itself falls back to the + naive plan rather than erroring. A scalar subquery always falls back, even + over an unqualified inner column, since that column resolves against the + subquery's own scope.) +- **Unknown table qualifier.** ``c.col`` when the only sides are ``a`` + and ``b``. +- **Unknown alias in GROUP BY / HAVING / ORDER BY.** A modifier-clause + column qualified with a table alias outside the join's two sides + (``c.col``) raises with a ``cannot resolve … in `` message + naming the expected aliases. +- **Right-side reference in SEMI / ANTI joins.** ``b.col`` or an + aggregate over ``b.col`` in the outer SELECT, GROUP BY, HAVING, ORDER + BY, or aggregate argument under ``SEMI JOIN`` / ``ANTI JOIN`` raises; + SEMI / ANTI return left-side columns only. (A right-side ``b.*`` falls + back to the naive-predicate plan with every other star — see the + soft-fallback list above.) +- **Aggregate of unqualified column.** ``SUM(score)`` raises; use + ``SUM(a.score)`` or ``SUM(b.score)``. (A qualified aggregate the IEJoin + cannot rebuild — ``SUM(a.score) OVER (...)``, ``SUM(a.score) FILTER + (WHERE ...)``, ``COUNT(*) * 2``, ``COUNT(a.*)`` — falls back to the + naive plan instead; see the soft-fallback list above.) +- **Catalog/schema-qualified extra predicates.** ``mycat.myschema.a.col`` + in an extra predicate raises; the alias rewriter would leave the + catalog/schema qualifier intact in the inner subquery, where it + references a different relation than intended. +- **Unqualified or unknown-aliased extra predicates.** ``WHERE strand + = '+'`` (unqualified) or ``WHERE c.score > 0`` (alias outside the + join's two sides) raise with an actionable message naming the + expected aliases. The naive-predicate plan would either silently bind the + column wrong or defer the error to the downstream engine, so the + dialect surfaces the user mistake at transpile time instead. + +The one argument-validation ``ValueError`` raise (it fires before any AST +inspection and is independent of the query shape): + +- **Unknown dialect string.** A dialect that resolves to neither a built-in + target (``None``, ``"duckdb"``, ``"datafusion"``) nor a custom target + registered on the plugin hub raises with the offending value echoed. + +Literal-range ``INTERSECTS`` (e.g. ``WHERE interval INTERSECTS +'chr1:1000-2000'``) is single-table and has no column-to-column join +to rewrite, so the dialect declines and the standard range-predicate +emission handles it identically to the ``dialect=None`` path. + DuckDB Optimizations ~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/transpilation/schema-mapping.rst b/docs/transpilation/schema-mapping.rst index 453c289..fd344b0 100644 --- a/docs/transpilation/schema-mapping.rst +++ b/docs/transpilation/schema-mapping.rst @@ -173,6 +173,36 @@ If your data uses 1-based coordinates (like VCF or GFF), configure the ], ) +.. note:: + + **Non-canonical encodings emit a capability-driven canonicalization wrapper.** + When a table declares an encoding other than the default 0-based half-open + (for example ``coordinate_system="1based"`` or ``interval_type="closed"``), + GIQL canonicalizes its coordinates by wrapping the relation in a hidden CTE. + The wrapper's projection form is chosen from the target's capabilities: the + ``"duckdb"`` target emits ``SELECT * REPLACE (...)`` (also supported by + BigQuery, Snowflake, and ClickHouse), while the generic (``dialect=None``) + and ``"datafusion"`` targets emit the portable ``SELECT * EXCEPT (start, end), + , `` form. The ``* EXCEPT`` form runs on ``* EXCEPT``-capable + engines (the DataFusion family) but is **not** SQL-92 and is **not** + DuckDB-runnable; it is row-equivalent to the ``* REPLACE`` form but re-appends + the recomputed interval columns at the end of the projection. Tables in the + default 0-based half-open encoding are unaffected -- they take an identity fast + path that emits portable SQL on every target. + + Neither form is SQL-92. To target a strict SQL-92 engine (PostgreSQL, SQLite), + store your data in 0-based half-open form, or convert it explicitly in a CTE + and reference that CTE (which GIQL treats as already canonical). Such a CTE -- + and any CTE or subquery passed as an operator reference -- must project the + canonical ``chrom`` / ``start`` / ``end`` columns; GIQL validates this contract + at transpile time and raises a ``ValueError`` naming the missing column(s) + rather than emitting SQL that fails with an engine ``column not found`` error. + + The projection form is chosen from the active target's + :class:`~giql.targets.Capabilities` (``supports_star_replace``). A custom + engine can select either form by registering a target with the appropriate + capabilities -- see :doc:`extending`. + Working with Point Features ~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/pyproject.toml b/pyproject.toml index 647358b..cc61041 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dev = [ "pytest>=7.0.0", "ruff>=0.1.0", "datafusion>=52.3.0", + "pytest-manifest>=1.0", ] docs = [ "sphinx>=7.0", @@ -57,6 +58,10 @@ path = "build-hooks/metadata.py" [tool.pytest.ini_options] addopts = "--cov --cov-config=.coveragerc" +markers = [ + "integration: end-to-end tests that execute generated SQL against a real database", + "usage: the canonical query-usage pattern (giql.tests usage_patterns) a test exercises", +] [tool.ruff] line-length = 89 @@ -89,3 +94,8 @@ pandas = ">=2.0.0" sqlglot = ">=20.0.0,<30" pip = "*" hypothesis = ">=6.148.2,<7" + +[tool.pixi.pypi-dependencies] +# pytest-manifest has no conda-forge package; the test suite's snapshot +# fixtures need it, so install it from PyPI into the pixi environment. +pytest-manifest = ">=1.0" diff --git a/src/giql/__init__.py b/src/giql/__init__.py index e5df351..60aced8 100644 --- a/src/giql/__init__.py +++ b/src/giql/__init__.py @@ -3,15 +3,36 @@ A SQL dialect for genomic range queries. """ -from giql.constants import DEFAULT_BIN_SIZE +from giql.expander import REGISTRY +from giql.expander import ExpanderRegistry +from giql.expander import ExpansionContext +from giql.expander import OperatorExpander +from giql.expander import StatementFinalizer +from giql.expander import register from giql.table import Table +from giql.targets import Capabilities +from giql.targets import DataFusionTarget +from giql.targets import DuckDBTarget +from giql.targets import GenericTarget +from giql.targets import Target from giql.transpile import transpile __version__ = "0.1.0" __all__ = [ - "DEFAULT_BIN_SIZE", "Table", "transpile", + # Extension hook: register custom targets and override operator expanders. + "register", + "REGISTRY", + "ExpanderRegistry", + "ExpansionContext", + "OperatorExpander", + "StatementFinalizer", + "Target", + "Capabilities", + "GenericTarget", + "DuckDBTarget", + "DataFusionTarget", ] diff --git a/src/giql/canonical.py b/src/giql/canonical.py new file mode 100644 index 0000000..f019a42 --- /dev/null +++ b/src/giql/canonical.py @@ -0,0 +1,91 @@ +"""Convert raw column expressions to and from canonical 0-based half-open coordinates. + +Genomic intervals can be expressed in 0-based half-open, 0-based closed, +1-based half-open, or 1-based closed coordinates. The transpiler emits SQL +in 0-based half-open by canonicalizing every column endpoint at the source. +These helpers wrap a raw column reference with the offset needed for that +canonicalization based on the source table's declared coordinate system. +The ``decanonical_*`` helpers apply the inverse offset to emit an endpoint +back in the table's own coordinate system. (Literal ranges go through +:meth:`giql.range_parser.RangeParser.to_zero_based_half_open`.) +""" + +from giql.table import Table + + +def canonical_start(raw_start: str, table: Table | None) -> str: + """Wrap a raw start column expression to canonical 0-based half-open start. + + When ``table`` is ``None``, the raw expression is returned unchanged + (treated as 0-based half-open). + + Conversion by :attr:`Table.coordinate_system`: + + - ``0based``: ``start`` (identity) + - ``1based``: ``start - 1`` + """ + if table is None or table.coordinate_system == "0based": + return raw_start + return f"({raw_start} - 1)" + + +def canonical_end(raw_end: str, table: Table | None) -> str: + """Wrap a raw end column expression to canonical 0-based half-open end. + + When ``table`` is ``None``, the raw expression is returned unchanged + (treated as 0-based half-open). + + Conversion by (:attr:`Table.coordinate_system`, :attr:`Table.interval_type`): + + - ``0based`` / ``half_open``: ``end`` (identity) + - ``0based`` / ``closed``: ``end + 1`` + - ``1based`` / ``half_open``: ``end - 1`` + - ``1based`` / ``closed``: ``end`` (identity) + """ + if table is None: + return raw_end + key = (table.coordinate_system, table.interval_type) + if key == ("0based", "closed"): + return f"({raw_end} + 1)" + if key == ("1based", "half_open"): + return f"({raw_end} - 1)" + return raw_end + + +def decanonical_start(canonical_expr: str, table: Table | None) -> str: + """Convert a canonical 0-based half-open start to the table's encoding. + + Inverse of :func:`canonical_start`. When ``table`` is ``None``, the + expression is returned unchanged (treated as 0-based half-open). + + Conversion by :attr:`Table.coordinate_system`: + + - ``0based``: ``start`` (identity) + - ``1based``: ``start + 1`` + """ + if table is None or table.coordinate_system == "0based": + return canonical_expr + return f"({canonical_expr} + 1)" + + +def decanonical_end(canonical_expr: str, table: Table | None) -> str: + """Convert a canonical 0-based half-open end to the table's encoding. + + Inverse of :func:`canonical_end`. When ``table`` is ``None``, the + expression is returned unchanged (treated as 0-based half-open). + + Conversion by (:attr:`Table.coordinate_system`, :attr:`Table.interval_type`): + + - ``0based`` / ``half_open``: ``end`` (identity) + - ``0based`` / ``closed``: ``end - 1`` + - ``1based`` / ``half_open``: ``end + 1`` + - ``1based`` / ``closed``: ``end`` (identity) + """ + if table is None: + return canonical_expr + key = (table.coordinate_system, table.interval_type) + if key == ("0based", "closed"): + return f"({canonical_expr} - 1)" + if key == ("1based", "half_open"): + return f"({canonical_expr} + 1)" + return canonical_expr diff --git a/src/giql/canonicalizer.py b/src/giql/canonicalizer.py new file mode 100644 index 0000000..f3e3061 --- /dev/null +++ b/src/giql/canonicalizer.py @@ -0,0 +1,575 @@ +"""Pass 2 of the GIQL normalization pipeline: ``CanonicalizeCoordinates``. + +This module implements the second normalization pass described in epic #114. It +runs after :func:`giql.resolver.resolve_operator_refs` (pass 1) and consumes the +:class:`~giql.resolver.ResolvedRef` metadata that pass attached to every GIQL +operator slot. For every interval-bearing relation reachable from an operator +slot whose ``ResolvedRef`` carries a :class:`~giql.table.Table` config declaring +a *non-canonical* encoding (``coordinate_system`` / ``interval_type`` other than +``0based`` / ``half_open``), the pass synthesizes a hidden ``__giql_canon_*`` +WITH-CTE that projects the relation to canonical 0-based half-open coordinates +(replicating the :mod:`giql.canonical` arithmetic as SQL in the CTE projection) +and rewrites the operator slot — both the AST node and its ``ResolvedRef`` +metadata — to point at the canonical CTE. + +The synthesis follows ``sqlglot``'s ``eliminate_subqueries`` mechanics +(:mod:`sqlglot.optimizer.eliminate_subqueries`): + +* **Collision-safe aliases.** Candidate names come from + :func:`sqlglot.helper.name_sequence` with the ``__giql_canon_`` prefix, checked + against a taken-set collected across every scope's table sources and CTE + aliases; :func:`sqlglot.helper.find_new_name` resolves any residual collision + against a user identifier. +* **Wrapper-body deduplication.** Two operator slots referencing the same source + under the same encoding share a single canonical CTE — one source is wrapped at + most once per query, keyed on the rendered wrapper body. +* **DAG-ordered insertion.** Canonical CTEs depend only on base (registered) + tables, never on other CTEs, so they are prepended to the outermost ``WITH``; + any existing CTE or main-body operator that now references a canonical CTE + therefore sees it defined first. + +Identity-encoded relations pass through **unwrapped**: a relation already in +0-based half-open form, or one with no ``Table`` config (a CTE or subquery +reference, assumed canonical), is left untouched — the readability-vs-volume +tradeoff the epic calls out (only synthesize a wrapper when canonicalization +actually changes columns). + +Engine portability (capability-driven, issue #145) +--------------------------------------------------- +The wrapper projection canonicalizes the interval columns while passing every +other source column through untouched (the registry declares only the genomic +columns, so an explicit full-column projection is not available). The emit +strategy is chosen from the active target's :class:`~giql.targets.Capabilities`, +following the DISJOIN passthrough's precedent (#143) — the same two emit forms, +though this wrapper additionally accepts ``capabilities=None`` (a direct caller +default, see :func:`canonicalize_coordinates`): + +* ``SELECT * REPLACE (...)`` when ``capabilities.supports_star_replace`` holds + (DuckDB / BigQuery / Snowflake / ClickHouse) — substitutes start/end in place, + preserving source column order; +* the portable ``SELECT * EXCEPT (start, end), , `` form otherwise + (the generic baseline / DataFusion family), which every ``* EXCEPT``-capable + engine plans. This form is row-equivalent but **not column-order-equivalent**: + ``* EXCEPT`` drops the interval columns and re-appends the recomputed ones at + the end of the projection. It is also not SQL-92 and not DuckDB-runnable. + +Identity-encoded (default 0-based half-open) relations are unaffected either way: +they skip wrapping entirely and emit portable SQL. The capability is threaded in +from :func:`giql.transpile.transpile` via the active target; a direct caller that +passes no capabilities defaults to the ``* REPLACE`` form (the historical +behavior). This finalizes the dialect-aware emit strategy formerly tracked by +https://github.com/abdenlab/giql/issues/132. + +Gating (epic #114, step 6) +-------------------------- +The pass is gated per operator by a ``GIQL_CANONICALIZE`` class attribute on the +operator's expression class. An operator opts in by setting +``GIQL_CANONICALIZE = True``; absent or ``False`` the pass ignores it entirely. +The operator port issues — #122 (DISJOIN) and #123 (NEAREST / DISTANCE / +predicates) — flipped these flags as each operator's emitter moved off in-emitter +canonicalization (:mod:`giql.canonical`) and onto this pass's output. As of those +ports every migrated operator opts in by default, so the pass actively +synthesizes wrappers; an operator can still toggle its flag off (a test or a +not-yet-migrated operator), in which case the pass leaves it untouched and the +emitted SQL is byte-identical for it. + +De-canonicalization hook +------------------------- +A migrated operator's *output* columns must land back in the target relation's +declared encoding. Epic #114 step 6 envisioned a rewrite of the outermost +``SELECT`` projection, but that placement is wrong for a table function: DISJOIN +and NEAREST synthesize their output columns and pass whole source rows through at +*generation* time, so those columns do not exist as AST in this pass, and a +``SELECT *`` consumer hides them from any outer-projection rewrite. So +:func:`_decanonicalize_outputs` instead records each wrapped slot's *original* +:class:`~giql.table.Table` on the operator's +:class:`~giql.resolver.OperatorResolution`, and the operator's emitter reads it +to de-canonicalize those synthesized columns where it generates them (DISJOIN, +issue #122; NEAREST's target row passthrough, issue #123). + +Column / interval operands (epic #114, step 8 / issue #123) +----------------------------------------------------------- +A reference slot (DISJOIN/NEAREST target, DISJOIN reference) *owns* the relation +it names, so the pass can wrap that whole relation in a ``__giql_canon_*`` CTE and +redirect the slot's AST node to it. A *column* operand cannot be wrapped that way: +``DISTANCE(a.interval, b.interval)`` and ``a.interval INTERSECTS b.interval`` +reference an alias bound in the *enclosing* query's ``FROM`` / ``JOIN``, shared +with the user's own projection (``SELECT a.start, DISTANCE(...)``). Rewriting that +``FROM`` source to a canonical CTE would silently canonicalize the user's own +``a.start`` too — a behavior change. NEAREST's column / implicit-outer +``reference`` slot is the same shape (an alias from the outer LATERAL relation). + +For those operands the pass therefore canonicalizes the resolution metadata *in +place* rather than synthesizing a CTE: :func:`_canonicalize_column_operands` +rewrites each :class:`~giql.resolver.ResolvedColumn` (DISTANCE's two operands and +the spatial predicates' column operands) and each non-table +:class:`~giql.resolver.ResolvedInterval` (NEAREST's ``column`` / +``implicit_outer`` reference) so its ``start`` / ``end`` fragments carry the +canonical 0-based half-open arithmetic and its ``table`` is blanked. The emitter +then consumes the fragments verbatim — no in-emitter +:func:`giql.canonical.canonical_start` / ``canonical_end``. The arithmetic is the +same the emitter used to emit inline, so the SQL stays byte-identical for these +operands; only the *owner* of the arithmetic moves from the generator to the pass. +A ``literal_range`` interval is already canonical and is left untouched, and an +operand whose ``table`` is already canonical (or ``None``) is a no-op. +""" + +from __future__ import annotations + +from dataclasses import replace + +from sqlglot import exp +from sqlglot.helper import find_new_name +from sqlglot.helper import name_sequence +from sqlglot.optimizer.scope import traverse_scope + +from giql.expressions import Contains +from giql.expressions import GIQLDisjoin +from giql.expressions import GIQLDistance +from giql.expressions import GIQLNearest +from giql.expressions import Intersects +from giql.expressions import SlotSpec +from giql.expressions import SpatialSetPredicate +from giql.expressions import Within +from giql.resolver import META_KEY +from giql.resolver import OperatorResolution +from giql.resolver import ResolvedColumn +from giql.resolver import ResolvedInterval +from giql.resolver import ResolvedRef +from giql.table import Table +from giql.targets import Capabilities + +__all__ = [ + "CANON_PREFIX", + "canonicalize_coordinates", +] + +#: The alias prefix for every synthesized canonical wrapper CTE. Mirrors the +#: ``__giql_dj_`` reserved prefix the DISJOIN emitter already uses; the leading +#: double underscore keeps the namespace clear of user identifiers. +CANON_PREFIX = "__giql_canon_" + +#: The GIQL operator expression classes the pass inspects — the same set pass 1 +#: annotates. Membership alone does not opt an operator in; the per-class +#: ``GIQL_CANONICALIZE`` flag does (see module docstring). +_OPERATORS: tuple[type[exp.Expression], ...] = ( + GIQLDisjoin, + GIQLNearest, + GIQLDistance, + Intersects, + Contains, + Within, + SpatialSetPredicate, +) + + +def canonicalize_coordinates( + expression: exp.Expression, capabilities: Capabilities | None = None +) -> exp.Expression: + """Synthesize canonical wrapper CTEs for non-canonical operator operands. + + Walks *expression* for opted-in GIQL operators (those whose expression class + sets ``GIQL_CANONICALIZE = True``), and for each reference slot that resolved + to a registered table with a non-canonical encoding, synthesizes a + ``__giql_canon_*`` WITH-CTE projecting that table to canonical 0-based + half-open coordinates and rewrites the slot (AST node + ``ResolvedRef`` + metadata) to point at the canonical CTE. + + The pass mutates and returns *expression* in place. For an operator whose + ``GIQL_CANONICALIZE`` flag is off, or whose operands are already in the + canonical 0-based half-open encoding, it touches nothing and leaves the + emitted SQL byte-identical. + + Parameters + ---------- + expression : exp.Expression + The pass-1-annotated AST. + capabilities : Capabilities | None + The active target's capabilities, used to choose the wrapper projection's + emit strategy (``* REPLACE`` vs the portable ``* EXCEPT`` form — see the + module docstring). :func:`giql.transpile.transpile` passes the active + target's capabilities; ``None`` (a direct caller) defaults to the + ``* REPLACE`` form, preserving the historical behavior. + + Returns + ------- + exp.Expression + The same *expression*, with canonical wrapper CTEs inserted and the + opted-in operator slots that reference non-canonical tables rewritten to + point at them. + """ + # Column / interval operands (DISTANCE, predicates, NEAREST's non-table + # reference) canonicalize their metadata in place; this is independent of the + # ref-slot CTE synthesis below and runs for every opted-in operator. + _canonicalize_column_operands(expression) + + targets = _collect_targets(expression) + if not targets: + # No opted-in operator carries a non-canonical *reference-slot* operand: + # no wrapper CTE is synthesized (the in-place column canonicalization + # above already ran). + return expression + + taken = _collect_taken_names(expression) + next_name = name_sequence(CANON_PREFIX) + body_to_name: dict[str, str] = {} + new_ctes: list[exp.CTE] = [] + + for node, arg, ref in targets: + body = _canonical_projection(ref, capabilities) + body_sql = body.sql() + name = body_to_name.get(body_sql) + if name is None: + # First time this exact wrapper body is needed: mint a fresh, + # collision-safe alias and emit one CTE for it. + name = _fresh_name(next_name, taken) + taken.add(name) + body_to_name[body_sql] = name + new_ctes.append(_make_cte(name, body)) + # Subsequent slots with an identical body reuse the same CTE (dedup). + _rewrite_slot(node, arg, ref, name) + + _insert_ctes(expression, new_ctes) + _decanonicalize_outputs(expression, targets) + return expression + + +def _collect_targets( + expression: exp.Expression, +) -> list[tuple[exp.Expression, str, ResolvedRef]]: + """Return ``(node, slot_arg, ref)`` triples needing canonicalization. + + A slot qualifies when its operator opts in (``GIQL_CANONICALIZE``), the slot + is a reference slot that resolved to a registered table, and that table's + declared encoding is non-canonical. CTE / subquery references (assumed + canonical) and identity-encoded tables are skipped — they pass through + unwrapped. + """ + targets: list[tuple[exp.Expression, str, ResolvedRef]] = [] + for node in expression.walk(): + if not isinstance(node, _OPERATORS): + continue + if not getattr(node, "GIQL_CANONICALIZE", False): + continue + resolution = node.meta.get(META_KEY) + if not isinstance(resolution, OperatorResolution): + continue + specs: tuple[SlotSpec, ...] = getattr(node, "GIQL_SLOTS", ()) + for spec in specs: + if not spec.is_ref_slot: + continue + ref = resolution.slots.get(spec.arg) + if ref is None: + continue + # Only registered tables carry a declared encoding; CTE / subquery + # references are assumed canonical and pass through unwrapped. + if ref.kind != "registered_table": + continue + if _is_canonical(ref.table): + continue + targets.append((node, spec.arg, ref)) + return targets + + +def _is_canonical(table: Table | None) -> bool: + """Whether *table*'s declared encoding is already canonical 0-based half-open. + + ``None`` (an unregistered relation) is treated as canonical, matching the + :mod:`giql.canonical` convention. + """ + if table is None: + return True + return table.coordinate_system == "0based" and table.interval_type == "half_open" + + +def _canonicalize_column_operands(expression: exp.Expression) -> None: + """Canonicalize column / interval operand metadata in place for opted-in ops. + + For every opted-in operator (``GIQL_CANONICALIZE``) carrying column operands + (DISTANCE's two operands and the spatial predicates' column operands, in the + :attr:`OperatorResolution.columns` channel) or a non-table interval reference + (NEAREST's ``column`` / ``implicit_outer`` slot, in + :attr:`OperatorResolution.slots`), rewrite the operand's ``start`` / ``end`` + SQL fragments to carry the canonical 0-based half-open arithmetic for its + declared encoding and blank its ``table``. + + This replaces the in-emitter :func:`giql.canonical.canonical_start` / + ``canonical_end`` wrapping for those operands (epic #114, step 8). Unlike the + ref-slot CTE synthesis, no relation is wrapped: the operand references an + alias bound in the enclosing query's ``FROM`` shared with the user's own + projection, so canonicalizing the whole relation would change unrelated + columns. Operands already canonical (``table`` is ``None`` or 0-based + half-open) are left untouched, keeping their SQL byte-identical; a + ``literal_range`` interval is already canonical and is skipped. + """ + for node in expression.walk(): + if not isinstance(node, _OPERATORS): + continue + if not getattr(node, "GIQL_CANONICALIZE", False): + continue + resolution = node.meta.get(META_KEY) + if not isinstance(resolution, OperatorResolution): + continue + for arg, column in list(resolution.columns.items()): + resolution.columns[arg] = _canonicalize_column(column) + for arg, slot in list(resolution.slots.items()): + if isinstance(slot, ResolvedInterval): + resolution.slots[arg] = _canonicalize_interval(slot) + + +def _canonicalize_column(column: ResolvedColumn) -> ResolvedColumn: + """Return *column* with canonical start/end fragments and a blanked table. + + A no-op (returns *column* unchanged) when its backing table is already + canonical or ``None``. + """ + if _is_canonical(column.table): + return column + return replace( + column, + start=_canonical_start_sql(column.start, column.table), + end=_canonical_end_sql(column.end, column.table), + table=None, + ) + + +def _canonicalize_interval(interval: ResolvedInterval) -> ResolvedInterval: + """Return *interval* with canonical start/end fragments and a blanked table. + + A no-op for a ``literal_range`` (already canonical, no table) and for any + interval whose backing table is already canonical or ``None``. + """ + if interval.kind == "literal_range" or _is_canonical(interval.table): + return interval + return replace( + interval, + start=_canonical_start_sql(interval.start, interval.table), + end=_canonical_end_sql(interval.end, interval.table), + table=None, + ) + + +def _canonical_start_sql(start: str, table: Table | None) -> str: + """SQL-fragment analog of :func:`giql.canonical.canonical_start`. + + - ``0based``: ``start`` (identity) + - ``1based``: ``(start - 1)`` + """ + if table is None or table.coordinate_system == "0based": + return start + return f"({start} - 1)" + + +def _canonical_end_sql(end: str, table: Table | None) -> str: + """SQL-fragment analog of :func:`giql.canonical.canonical_end`. + + - ``0based`` / ``half_open``: ``end`` (identity) + - ``0based`` / ``closed``: ``(end + 1)`` + - ``1based`` / ``half_open``: ``(end - 1)`` + - ``1based`` / ``closed``: ``end`` (identity) + """ + if table is None: + return end + key = (table.coordinate_system, table.interval_type) + if key == ("0based", "closed"): + return f"({end} + 1)" + if key == ("1based", "half_open"): + return f"({end} - 1)" + return end + + +def _collect_taken_names(expression: exp.Expression) -> set[str]: + """Collect every relation name already in use across all scopes. + + Mirrors ``eliminate_subqueries``'s taken-set: every registered-table source + name plus every CTE alias, so a freshly minted ``__giql_canon_*`` alias can + be proven collision-free. A scope-traversal failure falls back to a plain + tree walk so the set is never under-populated. + """ + taken: set[str] = set() + try: + scopes = traverse_scope(expression) + except Exception: + scopes = [] + for scope in scopes: + for source in scope.sources.values(): + if isinstance(source, exp.Table): + taken.add(source.name) + taken.update(scope.cte_sources) + # Belt-and-suspenders: a raw walk catches any name the scope machinery did + # not surface (e.g. when traverse_scope bailed on an exotic construct). + for table in expression.find_all(exp.Table): + taken.add(table.name) + for cte in expression.find_all(exp.CTE): + taken.add(cte.alias) + return taken + + +def _fresh_name(next_name, taken: set[str]) -> str: + """Return a ``__giql_canon_*`` alias not present in *taken*. + + Draws sequential candidates from the ``name_sequence`` generator and, on the + rare collision with a user identifier, defers to ``find_new_name`` for a + suffixed variant — the same two-step ``eliminate_subqueries`` uses. + """ + candidate = next_name() + if candidate in taken: + candidate = find_new_name(taken, candidate) + return candidate + + +def _canonical_projection( + ref: ResolvedRef, capabilities: Capabilities | None +) -> exp.Select: + """Build the ``SELECT`` body that projects *ref*'s table to canonical form. + + The projection is a **full-row passthrough**: ``SELECT *`` keeps every + physical column of the source relation, and only the two interval columns — + ``start`` / ``end``, under their original physical names — are rewritten with + the :mod:`giql.canonical` arithmetic for the table's declared encoding. + ``chrom`` and every non-interval column flow through the star untouched. + + The emit strategy is chosen from *capabilities* (issue #145), following the + precedent of :func:`giql.expanders.disjoin._disjoin_passthrough` — the same + two emit forms, with an added ``capabilities is None`` arm for direct callers + (the passthrough always receives a concrete ``ctx.capabilities``): + + * ``SELECT * REPLACE (...)`` when ``supports_star_replace`` holds (or no + capabilities are supplied) — substitutes the interval columns in place, + preserving source column order; + * the portable ``SELECT * EXCEPT (start, end), , `` form otherwise + — drops the interval columns from the star and re-appends them recomputed. + Row-equivalent but not column-order-equivalent, and not DuckDB-runnable. + + The full row (rather than a bare ``chrom`` / ``start`` / ``end`` triple) is + required by table-function operators whose final projection passes the whole + source row through — DISJOIN's ``SELECT t.*`` (#122) — and by their join-back + semantics, which key on the source's physical columns. A CTE / subquery + reference that only needs the canonical interval triple still reads those + three columns from the same wrapper. + """ + _chrom, start, end = ref.cols + table = ref.table + relation = ref.name + # Quote the interval identifiers: the canonical column names are physical and + # routinely reserved words (the default genomic layout's ``start`` / ``end``), + # so the executed wrapper must quote them. + start_id = exp.to_identifier(start, quoted=True) + end_id = exp.to_identifier(end, quoted=True) + start_proj = exp.alias_(_canonical_start_expr(start, table), start_id) + end_proj = exp.alias_(_canonical_end_expr(end, table), end_id) + if capabilities is None or capabilities.supports_star_replace: + star = exp.Star(replace=[start_proj, end_proj]) + return exp.Select(expressions=[star]).from_(exp.to_table(relation)) + # Portable form: drop the interval columns from the star and re-project them + # recomputed under their own names. EXCEPT removes them from the row; the + # trailing projections add them back in canonical form. + star = exp.Star(except_=[exp.column(start_id), exp.column(end_id)]) + return exp.Select(expressions=[star, start_proj, end_proj]).from_( + exp.to_table(relation) + ) + + +def _canonical_start_expr(start: str, table: Table | None) -> exp.Expression: + """Canonical 0-based half-open start expression for a raw start column. + + SQL analog of :func:`giql.canonical.canonical_start`: + + - ``0based``: ``start`` (identity) + - ``1based``: ``start - 1`` + """ + col = exp.column(exp.to_identifier(start, quoted=True)) + if table is None or table.coordinate_system == "0based": + return col + return exp.paren(exp.Sub(this=col, expression=exp.Literal.number(1))) + + +def _canonical_end_expr(end: str, table: Table | None) -> exp.Expression: + """Canonical 0-based half-open end expression for a raw end column. + + SQL analog of :func:`giql.canonical.canonical_end`: + + - ``0based`` / ``half_open``: ``end`` (identity) + - ``0based`` / ``closed``: ``end + 1`` + - ``1based`` / ``half_open``: ``end - 1`` + - ``1based`` / ``closed``: ``end`` (identity) + """ + col = exp.column(exp.to_identifier(end, quoted=True)) + if table is None: + return col + key = (table.coordinate_system, table.interval_type) + if key == ("0based", "closed"): + return exp.paren(exp.Add(this=col, expression=exp.Literal.number(1))) + if key == ("1based", "half_open"): + return exp.paren(exp.Sub(this=col, expression=exp.Literal.number(1))) + return col + + +def _make_cte(name: str, body: exp.Select) -> exp.CTE: + """Wrap a projection *body* as a named CTE.""" + return exp.CTE( + this=body, + alias=exp.TableAlias(this=exp.to_identifier(name)), + ) + + +def _rewrite_slot(node: exp.Expression, arg: str, ref: ResolvedRef, name: str) -> None: + """Point operator slot *arg* at the canonical CTE *name*. + + Rewrites both the AST slot (to a bare table reference naming the CTE) and the + attached :class:`~giql.resolver.OperatorResolution` (a fresh + :class:`~giql.resolver.ResolvedRef` of kind ``"cte"`` carrying no ``Table`` + config — the CTE is canonical by construction). The physical column names are + preserved because the CTE re-exposes them under their original names. + """ + new_ref = replace(ref, kind="cte", name=name, table=None) + resolution = node.meta[META_KEY] + resolution.slots[arg] = new_ref + node.set(arg, exp.to_table(name)) + + +def _insert_ctes(expression: exp.Expression, new_ctes: list[exp.CTE]) -> None: + """Prepend *new_ctes* to the outermost ``WITH`` in DAG order. + + Canonical CTEs depend only on base tables, so prepending them guarantees any + existing CTE or main-body reference resolves them. An existing ``WITH`` is + extended in place; otherwise a fresh one is attached to the root query. + """ + if not new_ctes: + return + query = expression.expression if isinstance(expression, exp.DDL) else expression + existing = query.args.get("with_") + if existing is not None: + existing.set("expressions", new_ctes + list(existing.expressions)) + else: + query.set("with_", exp.With(expressions=new_ctes)) + + +def _decanonicalize_outputs( + expression: exp.Expression, + targets: list[tuple[exp.Expression, str, ResolvedRef]], +) -> None: + """Preserve each wrapped slot's original encoding for the emitter's output. + + A wrapped slot's :class:`~giql.resolver.ResolvedRef` is rewritten to a + ``Table``-free canonical-CTE ref, which would otherwise lose the + (non-canonical) encoding the operator's *output* must round-trip back into. + + The de-canonicalization itself cannot be applied on the AST in this pass for + a table-function operator: DISJOIN synthesizes its ``disjoin_*`` columns and + its passed-through interval at *generation* time, so those columns do not + exist as AST here, and a ``SELECT *`` consumer hides them from any + outer-projection rewrite. The originally-envisioned outermost-projection + rewrite (epic #114, step 6) is therefore wrong for projected + table-function columns; instead this hook records the per-slot original + :class:`~giql.table.Table` on the :class:`~giql.resolver.OperatorResolution`, + and the operator's emitter reads it to de-canonicalize those synthesized + columns where it generates them (see :issue:`122`). + + *targets* carries the original (pre-rewrite) refs, so ``ref.table`` is the + source relation's declared encoding. + """ + for node, arg, ref in targets: + resolution = node.meta.get(META_KEY) + if isinstance(resolution, OperatorResolution): + resolution.output_tables[arg] = ref.table diff --git a/src/giql/constants.py b/src/giql/constants.py index 0daf016..c1ea6ec 100644 --- a/src/giql/constants.py +++ b/src/giql/constants.py @@ -10,5 +10,34 @@ DEFAULT_STRAND_COL = "strand" DEFAULT_GENOMIC_COL = "interval" -# Default bin size for INTERSECTS binned equi-join optimization -DEFAULT_BIN_SIZE = 10_000 +#: The reserved alias prefix naming DISJOIN's internal CTEs (``__giql_dj_ref`` / +#: ``__giql_dj_tgt`` / ``__giql_dj_cuts`` / ...). Follows the same +#: reserved-prefix scheme as :data:`giql.canonicalizer.CANON_PREFIX`; the leading +#: double underscore keeps +#: the namespace clear of user identifiers. A single source of truth so the +#: resolver's reserved-prefix guard and the DISJOIN expander's CTE names and +#: guard cannot drift apart. +DJ_PREFIX = "__giql_dj_" + +#: The reserved column name for CLUSTER's synthesized per-row "new cluster" flag +#: (``CASE ... END AS __giql_is_new_cluster``), consumed by the outer +#: ``SUM(...) OVER (...)`` window that assigns cluster ids. The reserved ``__giql_`` +#: prefix keeps it clear of user columns: a user column named ``is_new_cluster`` +#: coexisting with ``SELECT *`` otherwise mis-binds the SUM (#161). A single source +#: of truth so the alias and its ``SUM`` consumer cannot drift apart. +CLUSTER_FLAG_COL = "__giql_is_new_cluster" + +#: The reserved column name for the cluster id MERGE synthesizes +#: (``CLUSTER(...) AS __giql_cluster_id``) and then groups by. Reserved-prefixed for +#: the same reason as :data:`CLUSTER_FLAG_COL`; a single source of truth so the alias +#: and its ``GROUP BY`` consumer cannot drift apart. +CLUSTER_ID_COL = "__giql_cluster_id" + +#: The reserved alias prefix CLUSTER materializes a star's aliased sibling items +#: under, inside the ``__giql_lag_calc`` subquery, so the outer star can EXCEPT them +#: (and the outer projection reference them) without colliding with a same-named base +#: column or another sibling. Reserved-prefixed for the same reason as +#: :data:`CLUSTER_FLAG_COL`; ``SELECT *, expr AS score`` would otherwise strip the base +#: ``score`` from the star, and ``SELECT *, 1 AS x, 2 AS x`` would duplicate the EXCEPT +#: entry (#190). +CLUSTER_SIBLING_PREFIX = "__giql_sibling_" diff --git a/src/giql/dialect.py b/src/giql/dialect.py index 6c70104..9098c51 100644 --- a/src/giql/dialect.py +++ b/src/giql/dialect.py @@ -13,6 +13,7 @@ from giql.expressions import Contains from giql.expressions import GIQLCluster +from giql.expressions import GIQLDisjoin from giql.expressions import GIQLDistance from giql.expressions import GIQLMerge from giql.expressions import GIQLNearest @@ -57,6 +58,7 @@ class Parser(Parser): "MERGE": GIQLMerge.from_arg_list, "DISTANCE": GIQLDistance.from_arg_list, "NEAREST": GIQLNearest.from_arg_list, + "DISJOIN": GIQLDisjoin.from_arg_list, } def _parse_comparison(self): diff --git a/src/giql/expander.py b/src/giql/expander.py new file mode 100644 index 0000000..0499ca8 --- /dev/null +++ b/src/giql/expander.py @@ -0,0 +1,728 @@ +"""The operator-expander registry and the ``ExpandOperators`` pass (epic #137). + +This module provides the dispatch infrastructure every GIQL operator expands +through (epic #137, complete): + +* an :class:`OperatorExpander` protocol — ``expand(node, ctx) -> exp.Expression``, + the unit that turns one GIQL operator node into standard sqlglot AST for the + active target; +* an :class:`ExpansionContext` carrying everything an expander needs — the pass-1 + :class:`~giql.resolver.OperatorResolution` metadata, the active + :class:`~giql.targets.Target` and its :class:`~giql.targets.Capabilities`, + collision-safe alias minting, and the registered :class:`~giql.table.Tables`; +* a :class:`ExpanderRegistry` keyed by ``(target, operator_type)`` resolving an + expander through the fallback chain ``(target, op)`` → ``(generic, op)``; +* a :func:`register` decorator — the **public extension hook** by which a user + adds a target or overrides an operator for one, e.g. + ``@register(DuckDBTarget, GIQLDisjoin)``; +* the registry's teardown seam — :meth:`ExpanderRegistry.unregister`, + :meth:`ExpanderRegistry.clear`, and ``in`` membership (``__contains__``) — the + **supported public API** for undoing a custom registration. A user who + registers a custom expander (e.g. in a test fixture or a plugin's teardown) + resets the process-wide :data:`REGISTRY` through these rather than reaching + into private state: ``REGISTRY.unregister(DuckDBTarget(), GIQLDisjoin)`` drops + one entry, ``REGISTRY.clear()`` drops them all, and + ``(DuckDBTarget(), GIQLDisjoin) in REGISTRY`` tests for one; +* the :class:`ExpandOperators` pass, which runs after + :func:`giql.canonicalizer.canonicalize_coordinates`, walks the AST, and + replaces every GIQL operator node with the registry's expansion. + +Dispatch +-------- +With every operator migrated (epic #137 complete), the pass rewrites *every* GIQL +operator node: it resolves an expander for ``(active target, operator type)`` +through the registry's fallback chain and replaces the node with the AST that +expander returns. The built-in expanders register at import time via +:mod:`giql.expanders`, so a built-in operator always resolves at least a +``(generic, op)`` expander; a resolution miss is therefore an internal invariant +violation (an operator with no registered expander) and the pass raises rather +than leaving an un-serializable GIQL node in the tree. There is no per-operator +opt-in flag and no legacy ``*_sql`` fallback — both were migration scaffolding +removed once every operator was migrated (#146). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable +from typing import Protocol +from typing import runtime_checkable + +from sqlglot import exp +from sqlglot.helper import name_sequence + +from giql.expressions import Contains +from giql.expressions import GIQLCluster +from giql.expressions import GIQLDisjoin +from giql.expressions import GIQLDistance +from giql.expressions import GIQLMerge +from giql.expressions import GIQLNearest +from giql.expressions import Intersects +from giql.expressions import SpatialSetPredicate +from giql.expressions import Within +from giql.resolver import META_KEY +from giql.resolver import OperatorResolution +from giql.resolver import ResolutionError +from giql.table import Tables +from giql.targets import GenericTarget +from giql.targets import Target + +__all__ = [ + "EXPAND_ALIAS_PREFIX", + "ExpansionContext", + "OperatorExpander", + "StatementFinalizer", + "ExpanderRegistry", + "RegistrySnapshot", + "REGISTRY", + "register", + "expand_operators", + "ExpandOperators", +] + +#: The prefix every alias minted by :meth:`ExpansionContext.alias` carries. The +#: leading double underscore keeps synthesized names clear of user identifiers, +#: mirroring the ``__giql_canon_`` / ``__giql_dj_`` reserved prefixes. +EXPAND_ALIAS_PREFIX = "__giql_x_" + + +class ExpansionContext: + """Everything an :class:`OperatorExpander` needs to expand one operator. + + A fresh context is built per operator node by :class:`ExpandOperators`. It + bundles the node's pass-1 resolution metadata together with the active + target, collision-safe alias minting, and the registered tables, so an + expander never has to reach back into the pass or re-resolve anything. + + Attributes + ---------- + node : exp.Expression + The GIQL operator node being expanded. + resolution : OperatorResolution | None + The metadata :func:`giql.resolver.resolve_operator_refs` attached to + *node* (its resolved reference slots, column operands, deferrals, and + de-canonicalization output tables). ``None`` only if the node was never + annotated — an internal invariant violation an expander may assert on. + target : Target + The active :class:`~giql.targets.Target`, exposing its + ``capabilities`` and ``sqlglot_dialect``. + tables : Tables + The registered :class:`~giql.table.Tables` container. + registry : ExpanderRegistry | None + The registry the pass is resolving against. Carried so a whole-query + rewrite (CLUSTER / MERGE) can re-enter :func:`expand_operators` over the + SELECT it just restructured and expand sibling operators it copied into + it, honoring a custom-registry pass run. ``None`` for a standalone + context built outside the pass. + + A node-local expander that needs to rewrite the *enclosing* statement (rather + than just replace its own node) registers a :data:`StatementFinalizer` via + :meth:`add_statement_finalizer`; the pass applies every finalizer to the + statement after all node-local replacements. This is the query-level seam for + a target whose expansion must reshape the enclosing statement — for example to + project internal helper columns out of a surfacing ``SELECT *``. + """ + + __slots__ = ( + "node", + "resolution", + "target", + "tables", + "registry", + "_alias_seq", + "_finalizers", + ) + + def __init__( + self, + node: exp.Expression, + resolution: OperatorResolution | None, + target: Target, + tables: Tables, + alias_seq: Callable[[], str] | None = None, + registry: ExpanderRegistry | None = None, + finalizers: list[StatementFinalizer] | None = None, + ) -> None: + self.node = node + self.resolution = resolution + self.target = target + self.tables = tables + self.registry = registry + # A single sequence is threaded across every context built for one + # ``ExpandOperators`` run so aliases minted for sibling operators never + # collide; a standalone context falls back to its own sequence. + self._alias_seq = alias_seq or name_sequence(EXPAND_ALIAS_PREFIX) + # A single finalizer list is likewise shared across one run's contexts so + # a finalizer registered while expanding one node is applied once, after + # every node-local replacement; a standalone context gets its own (inert + # unless someone drives it manually). + self._finalizers = finalizers if finalizers is not None else [] + + @property + def capabilities(self): + """The active target's :class:`~giql.targets.Capabilities`.""" + return self.target.capabilities + + def add_statement_finalizer(self, finalizer: StatementFinalizer) -> None: + """Register a query-level :data:`StatementFinalizer` for this run. + + The **query-level seam**: an expander is node-local — its return value + replaces only its own node — so a target that must rewrite the *enclosing* + statement (for example to project internal helper columns out of a + surfacing ``SELECT *``) registers a finalizer here instead. + :func:`expand_operators` applies every registered finalizer to the + statement, in registration order, *after* all node-local replacements + complete; each receives the current statement root and returns the + (possibly new) root. + + A finalizer's returned root is emitted verbatim — beyond a type check that + it is an :class:`~sqlglot.expressions.Expression`, it is **not** + semantically re-validated — so it must not reference columns or relations + absent from the projection it rewrites: a wrapper over an absent column + builds without error but fails at engine runtime. Finalizers registered on + a standalone context (one built outside the pass) are collected but never + applied. + """ + self._finalizers.append(finalizer) + + def alias(self) -> str: + """Mint a fresh, query-unique alias with the reserved expander prefix. + + Draws sequential names from a per-run :func:`sqlglot.helper.name_sequence` + so two expanders (or two slots of one expander) never collide. + """ + return self._alias_seq() + + +@runtime_checkable +class OperatorExpander(Protocol): + """Turns one GIQL operator node into standard sqlglot AST for a target. + + An expander is any callable-bearing object whose :meth:`expand` takes the + operator node and its :class:`ExpansionContext` and returns the sqlglot AST + the node expands to for the active target. The returned expression replaces + the operator node in the tree; the stock per-target serializer renders it. + + The protocol is :func:`~typing.runtime_checkable`, so the registry can assert + a registered object satisfies it. A plain function is *not* an + ``OperatorExpander`` (it has no ``expand`` method); register one by wrapping + it (see :func:`register`, which accepts either form). + + An expander's **return value** is node-local: ``expand(node, ctx) -> + exp.Expression`` returns the one expression that replaces the operator node in + place. When a target additionally needs to rewrite the *enclosing* statement — + for example to project internal helper columns away from a surfacing + ``SELECT *`` — the expander registers a query-level :data:`StatementFinalizer` + via :meth:`ExpansionContext.add_statement_finalizer`, applied to the statement + after every node-local replacement. The DuckDB INTERSECTS IEJoin whole-query + fold uses exactly this seam: its ``(DuckDBTarget, Intersects)`` expander + (:mod:`giql.expanders.intersects_duckdb`, #169) registers a finalizer that + replaces the statement root with the per-chromosome dynamic-SQL rewrite. + """ + + def expand(self, node: exp.Expression, ctx: ExpansionContext) -> exp.Expression: ... + + +#: A bare expander function — the form most expanders are written in. The +#: registry stores either an :class:`OperatorExpander` object or one of these. +ExpanderFn = Callable[[exp.Expression, ExpansionContext], exp.Expression] + +#: A query-level statement finalizer: ``finalize(root) -> root``. An expander +#: registers one via :meth:`ExpansionContext.add_statement_finalizer` to wrap or +#: rewrite the enclosing statement after every node-local replacement; the pass +#: applies each in registration order and threads the (possibly new) root through. +#: Used, for example, to project internal helper columns out of a surfacing +#: ``SELECT *`` / ``b.*``. +StatementFinalizer = Callable[[exp.Expression], exp.Expression] + + +def _as_callable(expander: OperatorExpander | ExpanderFn) -> ExpanderFn: + """Normalize an expander to a plain ``(node, ctx) -> Expression`` callable.""" + # ``@runtime_checkable`` only checks an ``expand`` attribute *exists*, so an + # object with a non-callable ``expand`` would pass the isinstance check; guard + # that the attribute is actually callable before trusting the protocol branch. + if isinstance(expander, OperatorExpander) and callable(expander.expand): + return expander.expand + if callable(expander): + return expander + raise TypeError( + "An operator expander must be an OperatorExpander (an object with an " + f"'expand' method) or a callable; got {type(expander).__name__}." + ) + + +@dataclass(frozen=True) +class RegistrySnapshot: + """An opaque save/restore token for :class:`ExpanderRegistry` state. + + Bundles the ``(target, operator)`` expander map and the declared-target map + captured by :meth:`ExpanderRegistry.snapshot`. Hand it back to + :meth:`ExpanderRegistry.restore` to return the registry to that state, or + compare two snapshots for equality to assert the registry returned to a + baseline (the leak-guard use). + + A snapshot is compared by equality and handed to :meth:`restore` only; it is + **not** hashable (its two ``dict`` fields are mutable), so ``hash(snapshot)`` + raises — never use one as a set member or dict key. + """ + + expanders: dict[tuple[Target, type], ExpanderFn] + targets: dict[str, Target] + + +class ExpanderRegistry: + """Maps ``(target, operator_type)`` to the expander that handles it. + + Resolution follows the epic's fallback chain so a generic expander written + once covers every target, and a per-target entry overrides it only where the + engine genuinely differs: + + 1. ``(target, op)`` — an expander registered for the *exact* active target; + 2. ``(generic, op)`` — the portable expander registered for + :class:`~giql.targets.GenericTarget`; + 3. ``None`` — no expander resolves. Every built-in operator registers a + ``(generic, op)`` expander, so this only arises for a cleared registry or + an operator with no generic fallback; the pass raises rather than emitting + (there is no legacy ``*_sql`` fallback). + + The registry keys on the frozen, value-equal :class:`~giql.targets.Target` + *instance* (two ``DuckDBTarget()`` compare and hash alike), so registering + against ``DuckDBTarget`` and resolving against a freshly built + ``DuckDBTarget()`` hit the same entry. Registering against the *class* + ``DuckDBTarget`` is a convenience the decorator instantiates. + """ + + def __init__(self) -> None: + self._expanders: dict[tuple[Target, type], ExpanderFn] = {} + self._targets_by_name: dict[str, Target] = {} + + def register( + self, + target: Target, + operator: type, + expander: OperatorExpander | ExpanderFn, + ) -> None: + """Register *expander* for ``(target, operator)``, overriding any prior. + + Parameters + ---------- + target : Target + The target instance the expander handles. Use + :class:`~giql.targets.GenericTarget` for the portable fallback. + operator : type + The GIQL operator expression class (e.g. + :class:`~giql.expressions.GIQLDisjoin`). + expander : OperatorExpander | ExpanderFn + The expander object or function. A later registration for the same + key replaces an earlier one (last-write-wins override). + + Notes + ----- + Registering an expander also declares *target* by name (see + :meth:`register_target`), so a custom target becomes selectable via + ``transpile(dialect=target.name)`` as soon as any expander is registered + against it — a user overriding one operator need not also declare the + target separately. + + A later registration for the same ``(target, operator)`` key replaces the + earlier one (last-write-wins), so a user override of a *built-in* + target-specific expander — notably the DuckDB IEJoin + ``(DuckDBTarget, Intersects)`` join rewrite + (:mod:`giql.expanders.intersects_duckdb`) — simply supersedes it in the + registry; the pass then dispatches the operator node to the user's + expander instead. (:meth:`has_override` reports whether such an exact + non-generic entry exists for introspection.) + """ + self._expanders[(target, operator)] = _as_callable(expander) + self.register_target(target) + + def register_target(self, target: Target) -> None: + """Declare *target* so ``transpile(dialect=target.name)`` resolves it. + + The registry doubles as the target plugin hub: a custom + :class:`~giql.targets.Target` registered here (directly, or as a side + effect of :meth:`register`) is resolvable by name through + :func:`giql.targets.resolve_target`. This is the path for a + *capability-only* target — one that overrides no operators (so it never + calls :meth:`register`) but declares a distinct + :class:`~giql.targets.Capabilities` set / ``sqlglot_dialect``. + + The built-in :class:`~giql.targets.GenericTarget` is intentionally **not** + declared: it is the portable fallback, selected only via ``dialect=None`` + (see :func:`giql.targets.resolve_target`), never by name. The skip is + *by value* (``target == GenericTarget()``), so registering an expander + against ``GenericTarget`` is a no-op here — but a *custom* + :class:`~giql.targets.Target` subclass is declared under its own ``name`` + even if that name happens to be ``"generic"`` (it does not compare equal to + the built-in). A later registration for the same name replaces an earlier + one. + """ + if target == GenericTarget(): + return + self._targets_by_name[target.name] = target + + def target(self, name: str) -> Target | None: + """Return the registered custom target named *name*, or ``None``. + + Consulted by :func:`giql.targets.resolve_target` for any dialect name + that is not a built-in, so a registered custom target is selectable via + ``transpile(dialect=name)``. + """ + return self._targets_by_name.get(name) + + def resolve(self, target: Target, operator: type) -> ExpanderFn | None: + """Return the expander for ``(target, operator)`` via the fallback chain. + + Tries the exact ``(target, op)`` entry, then the + ``(GenericTarget(), op)`` fallback, then ``None``. ``None`` means no + expander is registered; the :class:`ExpandOperators` pass treats that as + an internal invariant violation and raises (there is no legacy emitter). + + Resolution is the single dispatch path for every operator, including the + DuckDB IEJoin join rewrite: on ``duckdb`` this returns the built-in + ``(DuckDBTarget, Intersects)`` override + (:mod:`giql.expanders.intersects_duckdb`), and a user registration for the + same key replaces it (last-write-wins) so this returns the user's expander + instead (#169). There is no separate pre-pass for INTERSECTS anymore. + """ + fn = self._expanders.get((target, operator)) + if fn is not None: + return fn + if target != GenericTarget(): + fn = self._expanders.get((GenericTarget(), operator)) + if fn is not None: + return fn + return None + + def has_override(self, target: Target, operator: type) -> bool: + """Whether an exact non-generic ``(target, operator)`` entry is registered. + + Returns ``True`` only when *target* is not :class:`~giql.targets.GenericTarget` + and an exact ``(target, operator)`` entry is registered; the portable + ``(GenericTarget(), operator)`` fallback is *not* an override and does not + count here. + + Such an entry marks a target-specific override that supersedes the + ``(GenericTarget, operator)`` fallback — for example the built-in DuckDB + IEJoin ``(DuckDBTarget, Intersects)`` join rewrite + (:mod:`giql.expanders.intersects_duckdb`), or a plugin's own override of + it. This is an introspection helper (the dispatch itself flows through + :meth:`resolve`); it is *not* a gate on any pre-pass, since INTERSECTS join + strategy is fully registry-driven (#169) with no pre-pass to bypass. + """ + return target != GenericTarget() and (target, operator) in self._expanders + + def unregister(self, target: Target, operator: type) -> None: + """Drop the ``(target, operator)`` entry if present. + + Part of the registry's **public teardown seam**: a caller that + registered a custom expander (a plugin, or a test fixture) undoes one + registration with this rather than mutating private state. Resolving the + same key afterward falls back through the chain to ``None``. Dropping an + absent key is a no-op. + """ + self._expanders.pop((target, operator), None) + + def clear(self) -> None: + """Drop every registration — both expanders and declared targets. + + The bulk form of :meth:`unregister`, and the **public reset** for the + process-wide :data:`REGISTRY` — e.g. a test fixture that saves and + restores the registry around a body that registers custom expanders or + targets. + """ + self._expanders.clear() + self._targets_by_name.clear() + + def snapshot(self) -> RegistrySnapshot: + """Return a shallow copy of the current registrations and targets. + + The save half of a save/restore seam that supports test + baseline-isolation: capture the baseline with this and hand it back to + :meth:`restore` afterward, so the built-in registrations survive an + isolating fixture that would otherwise :meth:`clear` them permanently. + + The returned :class:`RegistrySnapshot` is a fresh, opaque value bundling + both the ``(target, operator)`` expander map and the declared-target map; + mutating the registry afterward does not affect it, and two snapshots + compare equal when the registry is in the same state. + """ + return RegistrySnapshot(dict(self._expanders), dict(self._targets_by_name)) + + def restore(self, snapshot: RegistrySnapshot) -> None: + """Replace all registrations with those captured by :meth:`snapshot`. + + The restore half of the save/restore seam that supports test + baseline-isolation. Drops every current entry and re-installs exactly the + *snapshot* contents (expanders and declared targets), so a fixture can + return the registry to a previously captured baseline regardless of what + its body registered or cleared. + """ + self._expanders.clear() + self._expanders.update(snapshot.expanders) + self._targets_by_name.clear() + self._targets_by_name.update(snapshot.targets) + + def __contains__(self, key: tuple[Target, type]) -> bool: + """Whether an *exact* ``(target, operator)`` entry is registered. + + Tests the exact key only — it does **not** walk the generic fallback + chain :meth:`resolve` follows (``(GenericTarget(), op) in registry`` is + the only way ``(target, op)`` reports present via the generic entry). + Part of the public teardown seam, for asserting a registration landed or + was torn down. + """ + return key in self._expanders + + def __len__(self) -> int: + """The number of registered ``(target, operator)`` entries. + + Part of the public introspection surface (alongside ``in``): emptiness is + observable without reaching into private state, so a teardown fixture or + leak guard can assert the registry is clear via ``len(registry)`` or + ``bool(registry)``. + """ + return len(self._expanders) + + def __bool__(self) -> bool: + """Whether any expander is registered (``False`` when empty).""" + return bool(self._expanders) + + +#: The process-wide registry the :func:`register` decorator writes to and the +#: :class:`ExpandOperators` pass reads from. The built-in expanders register into +#: it at import time via :mod:`giql.expanders`; the pass rewrites every GIQL +#: operator node by resolving its expander here. +REGISTRY = ExpanderRegistry() + + +def register( + target: type[Target] | Target, operator: type +) -> Callable[[OperatorExpander | ExpanderFn], OperatorExpander | ExpanderFn]: + """Register an operator expander for ``(target, operator)`` — the extension hook. + + The public decorator by which a built-in or user-supplied expander is added + to the process-wide :data:`REGISTRY`:: + + @register(DuckDBTarget, GIQLDisjoin) + def expand_disjoin(node, ctx): ... + + Parameters + ---------- + target : type[Target] | Target + The target the expander handles, given as the target *class* (the usual + form — it is instantiated with its default capabilities) or an explicit + instance. Pass :class:`~giql.targets.GenericTarget` for the portable + expander every target falls back to. + operator : type + The GIQL operator expression class the expander handles. + + Returns + ------- + Callable + A decorator that registers its argument and returns it unchanged, so the + decorated object stays usable on its own. + """ + if isinstance(target, type): + try: + target_instance: Target = target() + except TypeError as exc: + raise TypeError( + f"register() could not instantiate target class {target.__name__!r} " + "with no arguments: a custom Target subclass with required fields " + "cannot be defaulted. Pass an instance (e.g. " + f"@register({target.__name__}(...), ...)) or give its fields " + "defaults so the class form can construct it." + ) from exc + else: + target_instance = target + + def decorator( + expander: OperatorExpander | ExpanderFn, + ) -> OperatorExpander | ExpanderFn: + REGISTRY.register(target_instance, operator, expander) + return expander + + return decorator + + +# The GIQL operator expression classes the pass inspects. Every one is migrated, +# so each resolves an expander through the registry (see module docstring). +# Imported eagerly at module scope, as ``canonicalizer.py`` imports the same +# classes: ``giql.expressions`` does not import this module, so there is no cycle. +_GIQL_OPERATORS: tuple[type, ...] = ( + GIQLDisjoin, + GIQLNearest, + GIQLDistance, + GIQLCluster, + GIQLMerge, + Intersects, + Contains, + Within, + SpatialSetPredicate, +) + + +def expand_operators( + expression: exp.Expression, + target: Target, + tables: Tables, + registry: ExpanderRegistry | None = None, +) -> exp.Expression: + """Replace every GIQL operator node with its registry expansion. + + Pass 3 of the normalization pipeline (epic #137). Runs after + :func:`giql.canonicalizer.canonicalize_coordinates`. Every GIQL operator is + migrated, so the pass dispatches *every* operator node: it resolves an expander + for ``(target, operator type)`` through the registry's fallback chain and + replaces the node with the AST that expander returns. A resolution miss is an + internal invariant violation (a built-in operator always has at least a + ``(generic, op)`` expander) and raises — there is no legacy ``*_sql`` fallback. + + The pass mutates *expression* in place for the node-local replacements, then + applies any :data:`StatementFinalizer` an expander registered (via + :meth:`ExpansionContext.add_statement_finalizer`) to the statement in + registration order. A finalizer may return a *new* root, so callers must use + the return value rather than assume in-place mutation. + + Parameters + ---------- + expression : exp.Expression + The pass-1/pass-2-annotated AST. + target : Target + The active target whose expanders to dispatch to. + tables : Tables + The registered table configurations, threaded into each context. + registry : ExpanderRegistry | None + The registry to resolve against; defaults to the process-wide + :data:`REGISTRY`. + + Returns + ------- + exp.Expression + *expression* with each operator node replaced by its target-specific + expansion, after any registered statement finalizers have run — the same + object when no finalizer replaced the root, otherwise the finalized root. + """ + reg = registry if registry is not None else REGISTRY + operators = _GIQL_OPERATORS + alias_seq = name_sequence(EXPAND_ALIAS_PREFIX) + # Shared across every context this run builds: an expander that must rewrite + # the enclosing statement appends a finalizer here, applied after the walk. + finalizers: list[StatementFinalizer] = [] + + # Collect first, then mutate: replacing nodes mid-walk is unsafe. + pending: list[tuple[exp.Expression, ExpanderFn]] = [] + for node in expression.walk(): + if not isinstance(node, operators): + continue + fn = reg.resolve(target, type(node)) + if fn is None: + # Every built-in operator registers at least a ``(generic, op)`` + # expander at import time, so a miss means the registry was cleared or + # a custom target shadowed an operator without providing one — an + # internal/config error, not user input. Fail loudly rather than + # leaving an un-serializable GIQL node for the stock serializer (there + # is no legacy ``*_sql`` fallback anymore). + raise ValueError( + f"No expander registered for {type(node).__name__} on target " + f"{target.name!r}; the ExpandOperators pass cannot leave a GIQL " + "operator node un-expanded." + ) + pending.append((node, fn)) + + # Replace deepest-first: replacing an ancestor node detaches any collected + # descendant, whose later ``node.replace`` would be a silent no-op (the + # detached node's parent is gone), so the inner operator would never expand. + # Mutating leaves first keeps every still-collected ancestor attached. A + # detached node is skipped defensively as a second line of defence (#154). + pending.sort(key=lambda item: _node_depth(item[0]), reverse=True) + + for node, fn in pending: + if node is not expression and node.parent is None: + # An ancestor expansion already detached this node from the tree; + # replacing it would be a no-op, so skip it. + continue + resolution = node.meta.get(META_KEY) + if not isinstance(resolution, OperatorResolution): + # Pass 1 guarantees every operator node carries its resolution; a + # missing or malformed one is an internal invariant violation, not a + # user error, so fail loudly rather than minting a None-resolution + # context the expander would then dereference blindly. + raise ResolutionError( + f"{type(node).__name__} reached the ExpandOperators pass without " + "valid resolution metadata; pass 1 (resolve_operator_refs) must " + "run first and annotate every operator node." + ) + ctx = ExpansionContext( + node, + resolution, + target, + tables, + alias_seq, + registry=reg, + finalizers=finalizers, + ) + replacement = fn(node, ctx) + if not isinstance(replacement, exp.Expression): + raise TypeError( + f"expander {fn!r} for {type(node).__name__} returned " + f"{type(replacement).__name__}, not exp.Expression" + ) + if replacement is not node: + node.replace(replacement) + + # Apply any query-level finalizers an expander registered, in registration + # order, once every node-local replacement is in place. A finalizer receives + # the current root and returns the (possibly new) root to thread forward. + for finalize in finalizers: + expression = finalize(expression) + if not isinstance(expression, exp.Expression): + # Mirror the node-local return guard: a finalizer that forgets to + # return the root would otherwise make the pass silently return the + # non-Expression far from the cause. + raise TypeError( + f"statement finalizer {finalize!r} returned " + f"{type(expression).__name__}, not exp.Expression" + ) + + return expression + + +def _node_depth(node: exp.Expression) -> int: + """The number of ancestors above *node* (root is depth 0). + + Used to order the collect-then-replace loop deepest-first so replacing an + ancestor never detaches a still-pending descendant before its own replace. + """ + depth = 0 + parent = node.parent + while parent is not None: + depth += 1 + parent = parent.parent + return depth + + +class ExpandOperators: + """Callable wrapper giving :func:`expand_operators` a uniform pass shape. + + The transpile pipeline composes each normalization pass as a callable bound + to its :class:`~giql.table.Tables`; this wrapper gives the expansion pass the + same shape — construct it with the active target and tables, then call it on + the AST — so it slots into the pipeline uniformly after + :class:`giql.canonicalizer.canonicalize_coordinates`. + """ + + def __init__( + self, + target: Target, + tables: Tables, + registry: ExpanderRegistry | None = None, + ) -> None: + self.target = target + self.tables = tables + self.registry = registry + + def transform(self, expression: exp.Expression) -> exp.Expression: + """Run the expansion pass over *expression* and return it.""" + return expand_operators(expression, self.target, self.tables, self.registry) diff --git a/src/giql/expanders/__init__.py b/src/giql/expanders/__init__.py new file mode 100644 index 0000000..2628289 --- /dev/null +++ b/src/giql/expanders/__init__.py @@ -0,0 +1,37 @@ +"""Built-in operator expanders for epic #137. + +Importing this package registers every built-in expander as a side effect: +each submodule decorates its expander(s) with ``@register(...)`` at import +time, and this package imports all of them. The import is wired once (in +:mod:`giql.transpile`) so the process-wide ``REGISTRY`` is populated before the +first transpile. + +New operator modules are picked up automatically: drop a ``.py`` into +this package and it is imported here without editing this file. + +Modules whose name starts with ``_`` are skipped (private helpers, not +expanders). Submodules import in :func:`pkgutil.iter_modules` order, which sets +last-write-wins resolution-order precedence for overlapping registrations; an +import error here aborts the whole package import by design (a broken built-in +expander must not be silently skipped). +""" + +from __future__ import annotations + +import importlib +import pkgutil + +from giql.expander import REGISTRY + +for _module_info in pkgutil.iter_modules(__path__): + if _module_info.name.startswith("_"): + continue + importlib.import_module(f"{__name__}.{_module_info.name}") + +# Fail loudly if discovery registered nothing. Under zipimport or PEP-420 +# namespace-package layouts ``pkgutil.iter_modules`` can yield no submodules, +# silently leaving the registry unpopulated; assert at least one expander landed +# so zero-discovery surfaces here rather than as a mystery legacy-path +# fallthrough. The check is branch-agnostic: it names no one operator, so it +# holds in every wave-3 worktree regardless of which expanders ship. +assert len(REGISTRY) > 0, "giql.expanders auto-discovery registered no expanders" diff --git a/src/giql/expanders/_distance.py b/src/giql/expanders/_distance.py new file mode 100644 index 0000000..eb28f9f --- /dev/null +++ b/src/giql/expanders/_distance.py @@ -0,0 +1,117 @@ +"""Shared distance-CASE string builder for the NEAREST expander. + +The private ``_``-prefix keeps this module out of the ``expanders`` package +auto-discovery (see :mod:`giql.expanders`), so it is a plain helper module rather +than a registered expander. + +.. note:: + + KEEP IN SYNC: :func:`generate_distance_case` here and the AST builder in + :mod:`giql.expanders.distance` (``expand_distance``) produce the *same* + distance CASE by two routes. DISTANCE migrated to the AST expander (epic + #137, issue #140); this string form survives because NEAREST assembles its + ORDER BY / filter math as SQL fragments (string-first) and splices the CASE + into them. Any change to the distance math here must be mirrored in the + expander (and vice versa); the parity test in ``tests/test_distance_udf.py`` + guards drift. +""" + +from __future__ import annotations + + +def generate_distance_case( + chrom_a: str, + start_a: str, + end_a: str, + strand_a: str | None, + chrom_b: str, + start_b: str, + end_b: str, + strand_b: str | None, + stranded: bool = False, + signed: bool = False, +) -> str: + """Generate a SQL CASE expression for distance calculation. + + Distances follow bedtools ``closest -d`` semantics: overlapping + intervals report ``0``, book-ended (adjacent) intervals where + ``A.end == B.start`` in half-open coordinates report ``1``, and a raw + half-open gap of N bases reports ``N + 1``. The ``+ 1`` is applied to + the absolute gap magnitude before any directional sign, so a downstream + book-ended pair reports ``+1`` and an upstream one ``-1`` in signed mode. + + :param chrom_a: + Chromosome column for interval A + :param start_a: + Start column for interval A + :param end_a: + End column for interval A + :param strand_a: + Strand column for interval A (None if not stranded) + :param chrom_b: + Chromosome column for interval B + :param start_b: + Start column for interval B + :param end_b: + End column for interval B + :param strand_b: + Strand column for interval B (None if not stranded) + :param stranded: + Whether to use strand-aware distance calculation + :param signed: + Whether to return signed distance (negative for upstream, positive for + downstream) + :return: + SQL CASE expression + """ + if not stranded or strand_a is None or strand_b is None: + # Basic distance calculation without strand awareness + if signed: + # Signed distance: negative for upstream (B before A), + # positive for downstream (B after A) + return ( + f"CASE WHEN {chrom_a} != {chrom_b} THEN NULL " + f"WHEN {start_a} < {end_b} AND {end_a} > {start_b} THEN 0 " + f"WHEN {end_a} <= {start_b} THEN ({start_b} - {end_a} + 1) " + f"ELSE -({start_a} - {end_b} + 1) END" + ) + # Unsigned (absolute) distance + return ( + f"CASE WHEN {chrom_a} != {chrom_b} THEN NULL " + f"WHEN {start_a} < {end_b} AND {end_a} > {start_b} THEN 0 " + f"WHEN {end_a} <= {start_b} THEN ({start_b} - {end_a} + 1) " + f"ELSE ({start_a} - {end_b} + 1) END" + ) + + # Stranded distance calculation + # Return NULL if either strand is '.', '?', or NULL + # Calculate distance and multiply by -1 if first interval is on '-' strand + if signed: + # Stranded + signed: apply strand flip AND directional sign + return ( + f"CASE WHEN {chrom_a} != {chrom_b} THEN NULL " + f"WHEN {strand_a} IS NULL OR {strand_b} IS NULL THEN NULL " + f"WHEN {strand_a} = '.' OR {strand_a} = '?' THEN NULL " + f"WHEN {strand_b} = '.' OR {strand_b} = '?' THEN NULL " + f"WHEN {start_a} < {end_b} AND {end_a} > {start_b} THEN 0 " + f"WHEN {end_a} <= {start_b} THEN " + f"CASE WHEN {strand_a} = '-' THEN -({start_b} - {end_a} + 1) " + f"ELSE ({start_b} - {end_a} + 1) END " + f"ELSE " + f"CASE WHEN {strand_a} = '-' THEN ({start_a} - {end_b} + 1) " + f"ELSE -({start_a} - {end_b} + 1) END END" + ) + # Stranded but not signed: apply strand flip only + return ( + f"CASE WHEN {chrom_a} != {chrom_b} THEN NULL " + f"WHEN {strand_a} IS NULL OR {strand_b} IS NULL THEN NULL " + f"WHEN {strand_a} = '.' OR {strand_a} = '?' THEN NULL " + f"WHEN {strand_b} = '.' OR {strand_b} = '?' THEN NULL " + f"WHEN {start_a} < {end_b} AND {end_a} > {start_b} THEN 0 " + f"WHEN {end_a} <= {start_b} THEN " + f"CASE WHEN {strand_a} = '-' THEN -({start_b} - {end_a} + 1) " + f"ELSE ({start_b} - {end_a} + 1) END " + f"ELSE " + f"CASE WHEN {strand_a} = '-' THEN -({start_a} - {end_b} + 1) " + f"ELSE ({start_a} - {end_b} + 1) END END" + ) diff --git a/src/giql/expanders/_genomic.py b/src/giql/expanders/_genomic.py new file mode 100644 index 0000000..4ced34c --- /dev/null +++ b/src/giql/expanders/_genomic.py @@ -0,0 +1,422 @@ +"""Operator-neutral CLUSTER/MERGE expansion toolkit (epic #137, issue #163). + +The CLUSTER (:mod:`giql.expanders.cluster`) and MERGE +(:mod:`giql.expanders.merge`) expanders are both **whole-query rewrites** built +on the same generic plumbing: resolve the genomic columns of an enclosing FROM +source/relation, locate the operator in a top-level projection, guard the shapes +that have no coherent rewrite, and transplant a freshly-built SELECT onto the +original in place. None of that plumbing is specific to either operator. + +This module is that shared surface — "operator-neutral" in that no helper is +specific to one operator. The lone exception proves the rule: the CLUSTER/MERGE +co-occurrence guard (:func:`reject_cluster_merge_mix`) references both concrete +operator types, but it is *symmetric* between them rather than agnostic of them — +a guard owned by neither expander belongs here, not in either one. It was +extracted from ``cluster.py`` (#163) so the shared toolkit has an explicit home +rather than living in one expander and being imported by its sibling — an +arrangement that made the dependency direction between the two expanders look +accidental and invited further expander-to-expander imports. The extraction is a +pure code-move: the emitted SQL is unchanged. + +The name is ``_``-prefixed so the :mod:`giql.expanders` auto-discovery import +skips it — it registers no expander; it is a helper module the expanders import. +The one genuinely CLUSTER-specific reuse point, ``expand_cluster_query`` (MERGE +composes the CLUSTER rewrite to build its clustered subquery), deliberately stays +in ``cluster.py`` and is imported by ``merge.py`` as a documented composition. +""" + +from __future__ import annotations + +from typing import NamedTuple +from typing import TypeVar + +from sqlglot import exp + +from giql.constants import DEFAULT_CHROM_COL +from giql.constants import DEFAULT_END_COL +from giql.constants import DEFAULT_START_COL +from giql.constants import DEFAULT_STRAND_COL +from giql.expressions import GIQLCluster +from giql.expressions import GIQLMerge +from giql.table import Table +from giql.table import Tables + +_T = TypeVar("_T", bound=exp.Expression) + + +class GenomicColumns(NamedTuple): + """The resolved physical column names CLUSTER / MERGE operate over. + + Derived from the enclosing FROM table by :func:`genomic_columns`. A + :class:`~typing.NamedTuple` so it still unpacks and indexes positionally + (``chrom, start, end, strand = columns``) while giving the four fields names. + """ + + chrom: str + start: str + end: str + strand: str + + +_CANONICAL_COLUMNS = GenomicColumns( + DEFAULT_CHROM_COL, DEFAULT_START_COL, DEFAULT_END_COL, DEFAULT_STRAND_COL +) + + +def genomic_columns(select: exp.Select, tables: Tables) -> GenomicColumns: + """Return the ``(chrom, start, end, strand)`` columns for *select*'s FROM source. + + Part of the shared, operator-neutral CLUSTER/MERGE expansion toolkit. Resolves + the physical genomic columns of the enclosing FROM relation: a bare registered + table yields its configured mapping (and the default strand column when the table + declares none); a derived table or CTE is resolved *through* to the underlying + table when the subquery passes its columns through unchanged (``SELECT *`` / + ``SELECT alias.*``). An unregistered table, an empty FROM, or a projection that + already exposes the canonical ``chrom`` / ``start`` / ``end`` columns falls back + to those defaults. + + Unlike ``validate_projection_contracts`` / ``_resolve_disjoin_reference`` in + ``resolver.py``, which require canonical columns from an operator's reference + *operands*, this deliberately traces a star-passthrough FROM *source* to its + physical columns — CLUSTER/MERGE operate over the enclosing relation's own + coordinates, a different operator semantics. That FROM-source resolution lives in + the CLUSTER/MERGE toolkit rather than the pass-1 resolver because it serves the + whole-query rewrite's need to name the physical window/grouping columns, not the + resolver's reference-slot contract validation. + + :raises ValueError: When the FROM source is an untraceable relation — a derived + table, CTE, ``VALUES``, table-valued function, or ``LATERAL`` — whose genomic + columns cannot be determined (it neither exposes the canonical columns nor + resolves to a single registered mapping), or when the FROM combines multiple + sources through a join. Silently defaulting to the canonical columns in these + cases would emit SQL referencing columns that do not exist, so the ambiguity + is surfaced instead (#164). + """ + if select.args.get("joins"): + # Only the query being rewritten carries a join that ``transplant`` would + # drop; a join nested inside a derived-table/CTE body is preserved intact, + # so the rejection guards the top-level FROM, not the recursive descent. + raise ValueError( + "CLUSTER/MERGE over a join is not supported; run it over a single " + "relation, or over a subquery." + ) + mapping, opaque = _resolve_from_columns(select, tables, frozenset()) + if mapping is not None: + return mapping + if opaque: + raise ValueError( + "Cannot resolve the genomic columns for CLUSTER/MERGE over this FROM " + "clause: it is a derived table or CTE that neither exposes the canonical " + f"{DEFAULT_CHROM_COL!r} / {DEFAULT_START_COL!r} / {DEFAULT_END_COL!r} " + "columns nor resolves to a single registered table. Alias the genomic " + "columns to their canonical names inside the subquery, or run " + "CLUSTER/MERGE directly over a registered table." + ) + return _CANONICAL_COLUMNS + + +def _resolve_from_columns( + select: exp.Select, tables: Tables, seen: frozenset[str] +) -> tuple[GenomicColumns | None, bool]: + """Resolve *select*'s driving FROM source to a ``(mapping, opaque)`` result. + + *mapping* is the resolved :class:`GenomicColumns` when a registered mapping backs + the source, else ``None``. *opaque* is ``True`` when the source is a derived + table/CTE whose columns could not be traced — a signal the caller turns into a + hard error rather than a silent canonical default. + + Only the driving ``from_`` relation is resolved; any joins are ignored. A join on + the top-level rewrite target is rejected upstream in :func:`genomic_columns`, + while a join inside a traced-through derived-table/CTE body is preserved by the + rewrite, so its column list is taken from the driving relation as before (#164). + """ + from_clause = select.args.get("from_") + if from_clause is None: + return None, False + return _resolve_source(from_clause.this, select, tables, seen) + + +def _resolve_source( + source: exp.Expression, + enclosing: exp.Select, + tables: Tables, + seen: frozenset[str], +) -> tuple[GenomicColumns | None, bool]: + """Resolve one FROM/JOIN *source* node to a ``(mapping, opaque)`` result.""" + if isinstance(source, exp.Table): + name = source.name + cte_query = _find_cte(enclosing, name) + if cte_query is not None: + if name in seen: + # A CTE that (transitively) selects from itself: stop recursing and + # report it untraceable rather than looping forever. + return None, True + return _resolve_query(cte_query, tables, seen | {name}) + table = tables.get(name) + if table is not None: + return _mapping_for(table), False + # An unregistered bare table is assumed to expose the canonical columns, + # matching how the resolver treats unknown relations elsewhere. + return None, False + if isinstance(source, exp.Subquery): + return _resolve_query(source.this, tables, seen) + # A LATERAL, VALUES, or table-valued function has no traceable table mapping. + return None, True + + +def is_star_projection(expression: exp.Expression) -> bool: + """Return ``True`` for a bare ``*`` or a qualified ``rel.*`` projection. + + A bare star parses as an :class:`~sqlglot.expressions.Star`; a qualified ``rel.*`` + parses as an :class:`~sqlglot.expressions.Column` wrapping a ``Star``, not a bare + ``Star``. Both project every column of their source, so any caller that special-cases + "this projection surfaces all of the relation's columns" must recognize both shapes + identically. Defining the predicate once keeps the CLUSTER required-column dedup, the + CLUSTER outer-star rewrite, and this module's derived-table passthrough check from + drifting apart — a drift whose failure modes are exactly a dangling ``rel.*`` alias + and duplicated required columns (#185). + """ + return isinstance(expression, exp.Star) or ( + isinstance(expression, exp.Column) and isinstance(expression.this, exp.Star) + ) + + +def _resolve_query( + query: exp.Expression, tables: Tables, seen: frozenset[str] +) -> tuple[GenomicColumns | None, bool]: + """Resolve the columns a derived-table/CTE *query* exposes to ``(mapping, opaque)``. + + Any projection *containing* a star — a bare ``*``, a qualified ``alias.*``, or a + star alongside other items (``SELECT r.*, 1 AS extra``) — passes its FROM + source's columns through, so the mapping is resolved from that source. An explicit + projection instead names the output columns directly: it is transparent only when + it re-exposes the canonical genomic columns, otherwise its genomic columns cannot + be named and the source is reported opaque. + """ + if isinstance(query, exp.SetOperation): + # A set operation's arms share one column list; the left arm determines it. + return _resolve_query(query.left, tables, seen) + if not isinstance(query, exp.Select): + # Defensive: a parenthesized subquery / set-operation arm the grammar can + # reach is always a SELECT or a SetOperation (caught above, including + # INTERSECT / EXCEPT), so this total-function guard for any other body (e.g. a + # bare VALUES, which instead parses as an un-wrapped FROM source handled in + # _resolve_source) is not reachable via a GIQL query. + return None, True + + projections = query.expressions + passes_through = any(is_star_projection(projection) for projection in projections) + if passes_through: + return _resolve_from_columns(query, tables, seen) + + output_names = { + projection.alias if isinstance(projection, exp.Alias) else projection.name + for projection in projections + if isinstance(projection, (exp.Alias, exp.Column)) + } + canonical = {DEFAULT_CHROM_COL, DEFAULT_START_COL, DEFAULT_END_COL} + if canonical <= output_names: + return None, False + return None, True + + +def _find_cte(select: exp.Select, name: str) -> exp.Expression | None: + """Return the query body of an enclosing CTE named *name*, or ``None``.""" + node: exp.Expression | None = select + while node is not None: + with_clause = node.args.get("with_") + if with_clause is not None: + for cte in with_clause.expressions: + if cte.alias.lower() == name.lower(): + return cte.this + node = node.parent + return None + + +def _mapping_for(table: Table) -> GenomicColumns: + """Return the genomic column mapping configured on a registered *table*.""" + return GenomicColumns( + table.chrom_col, + table.start_col, + table.end_col, + table.strand_col or DEFAULT_STRAND_COL, + ) + + +def extract_stranded(stranded_expr: exp.Expression | None) -> bool: + """Coerce a CLUSTER/MERGE ``stranded`` operand to a bool. Shared toolkit. + + Mirrors the legacy per-transformer coercion exactly: a missing operand is + ``False``; an ``exp.Boolean`` yields its raw ``.this``; an ``exp.Literal`` + compares case-folded to ``TRUE``. The final two arms (``exp.Literal`` and the + string-truthiness fallback) are defensive — the GIQL grammar only ever produces + ``exp.Boolean`` for ``stranded := `` — retained for parity with the + legacy port. + """ + if stranded_expr is None: + return False + if isinstance(stranded_expr, exp.Boolean): + return stranded_expr.this + if isinstance(stranded_expr, exp.Literal): + return str(stranded_expr.this).upper() == "TRUE" + return str(stranded_expr).upper() in ("TRUE", "1", "YES") + + +def find_projected(select: exp.Select, op_type: type[_T]) -> list[_T]: + """Return *select*'s projected operators of *op_type* (bare or aliased). Toolkit. + + Shared CLUSTER/MERGE primitive: both expanders and the co-occurrence guard + locate their operator the same way — a top-level SELECT projection item that + either *is* the operator or is an ``exp.Alias`` wrapping it. + """ + found: list[_T] = [] + for expression in select.expressions: + if isinstance(expression, op_type): + found.append(expression) + elif isinstance(expression, exp.Alias) and isinstance(expression.this, op_type): + found.append(expression.this) + return found + + +def require_top_level_projection( + select: exp.Select, node: exp.Expression, op_type: type +) -> None: + """Raise if *node* is buried inside a projection expression. Shared toolkit. + + A CLUSTER / MERGE is only expandable as a *top-level* projection item — bare + or directly aliased — because the whole-query rewrite restructures the SELECT + around it. One nested inside a larger projection expression such as + ``ABS(CLUSTER(interval))`` has no coherent rewrite and would otherwise leak an + unexpanded operator to the generator, so fail loudly here (#144 A16). An + operator that parses *outside* the projection entirely (e.g. in WHERE / + ORDER BY) is not under any projection item and is left for the expander's + existing no-op path. + """ + operator = op_type.__name__.removeprefix("GIQL").upper() + for projection in select.expressions: + inner = projection.this if isinstance(projection, exp.Alias) else projection + if inner is node: + return + if any(descendant is node for descendant in inner.walk()): + raise ValueError( + f"{operator} must be a top-level projection item; it cannot be " + "nested inside another expression (e.g. a function call or " + "arithmetic)." + ) + + +def reject_cluster_merge_mix(select: exp.Select) -> None: + """Raise if *select* projects both a CLUSTER and a MERGE. Shared toolkit. + + The two are mutually incompatible in one SELECT: MERGE aggregates the rows + away (it rewrites the query into a ``GROUP BY`` over a clustered subquery) + while CLUSTER is a per-row window over those same rows, so no coherent single + query expresses both. The legacy pre-pass chained the two transformers and + emitted *non-executable* SQL for this shape (a window over ``GROUP BY``- + aggregated rows — a DuckDB ``BinderException``, never a leaked operator). The + new in-place ``transplant`` cannot express both at all: whichever expander + runs first rewrites the shared SELECT and strands the sibling as an unexpanded + node. Fail loudly here — mirroring the ``Multiple MERGE expressions not yet + supported`` guard — so the combination raises a clear diagnostic rather than + emitting broken SQL. + """ + if find_projected(select, GIQLCluster) and find_projected(select, GIQLMerge): + raise ValueError( + "CLUSTER and MERGE cannot be combined in a single SELECT; MERGE " + "aggregates rows while CLUSTER is a per-row window. Use them in " + "separate queries (e.g. CLUSTER over a subquery, or MERGE over one)." + ) + + +# Top-level ``exp.Select`` args the ``_transform_for_*`` helpers rebuild onto their +# freshly-built ``new`` query (the SELECT list, the FROM/WHERE/GROUP/HAVING that +# move into the inner subquery, and the ORDER that stays outer). ``joins`` is +# rejected upstream in :func:`genomic_columns` before any transplant runs. See +# :func:`transplant` for how everything outside this set and +# :data:`_PRESERVED_ROOT_ARGS` is handled. +_REBUILT_ROOT_ARGS = frozenset( + {"expressions", "from_", "where", "group", "having", "order", "joins"} +) + +# Outer clauses the ``_transform_for_*`` helpers do NOT rebuild onto ``new`` but +# that belong on the outer/final query the rewrite produces, so :func:`transplant` +# re-attaches each (see it for why dropping any corrupts the result). These are +# sqlglot ``exp.Select`` arg names — only ``with_`` carries a trailing underscore. +# (A named ``WINDOW`` clause is deliberately NOT here: CLUSTER duplicates the +# window-referencing projection into the inner subquery too, so an outer-only +# re-attach would still dangle; :func:`transplant` rejects it loudly instead.) +_PRESERVED_ROOT_ARGS = ("with_", "limit", "offset", "distinct", "qualify") + + +def transplant(select: exp.Select, new: exp.Select) -> None: + """Rewrite *select* into *new*, preserving *select*'s outer clauses + (:data:`_PRESERVED_ROOT_ARGS`) and failing loud on any it cannot carry, while + keeping *select*'s identity. + + Part of the shared CLUSTER/MERGE expansion toolkit. Clears every argument of + *select* and re-installs *new*'s, so *select* keeps its position in the + surrounding tree (and its identity as the object the pass returns) while taking + on the rewritten structure. This is how a whole-query rewrite is applied to a + *root* SELECT, which has no parent to ``replace`` through. + + Precondition: *new* MUST be a detached throwaway — a freshly-built + ``exp.Select``, as the ``_transform_for_*`` helpers return. Its children are + re-parented onto *select*, so passing a node still attached elsewhere would + corrupt that other tree. + + The ``_transform_for_*`` helpers build *new* from :data:`_REBUILT_ROOT_ARGS` + (the SELECT list plus FROM / WHERE / GROUP / HAVING / ORDER) but never *select*'s + other top-level clauses. The clauses in :data:`_PRESERVED_ROOT_ARGS` belong on + the outer query the rewrite produces, so dropping them corrupts the result: a + CLUSTER/MERGE over a CTE FROM would dangle the ``WITH``'s CTE reference the + rewritten ``__giql_lag_calc`` subquery still names (nested, for MERGE, inside + its ``__giql_clustered`` wrapper) — non-executable SQL (#174) — while a dropped + ``LIMIT`` / ``OFFSET`` / ``DISTINCT`` / ``QUALIFY`` would silently return the + wrong rows (#181). ``transplant`` re-attaches each onto the outer query. + + Any *other* top-level clause (e.g. a named ``WINDOW``, ``SELECT … INTO``, + ``FOR UPDATE``, ``SORT``/``CLUSTER``/``DISTRIBUTE BY``) is neither rebuilt nor + preserved, so rather than silently drop it — the very failure mode #174/#181 were + about — this raises :class:`ValueError`. Silent-drop is the worst outcome for a + query compiler; a loud error is diagnosable and future-proofs the rewrite against + new sqlglot args. (A named ``WINDOW`` referenced by a surviving projection is a + genuine gap rather than a nonsensical clause, but preserving it correctly means + threading the definition into the inner subquery too — deferred; for now it + fails loud rather than emitting a dangling ``OVER w``.) + + :raises ValueError: When *select* carries a top-level clause that is neither + rebuilt (:data:`_REBUILT_ROOT_ARGS`) nor preserved + (:data:`_PRESERVED_ROOT_ARGS`). + """ + assert new.parent is None, "transplant() requires a detached `new` subtree" + unhandled = sorted( + key + for key, value in select.args.items() + if value is not None + and key not in _REBUILT_ROOT_ARGS + and key not in _PRESERVED_ROOT_ARGS + ) + if unhandled: + raise ValueError( + "CLUSTER/MERGE cannot carry these top-level clause(s) through the " + f"whole-query rewrite: {unhandled}. They are neither rebuilt nor in the " + "preserved set (e.g. SELECT INTO, FOR UPDATE, SORT/CLUSTER/DISTRIBUTE " + "BY). Run the operator over a subquery that applies the clause instead." + ) + preserved = {key: select.args.get(key) for key in _PRESERVED_ROOT_ARGS} + select.args.clear() + for key, value in list(new.args.items()): + select.set(key, value) + for key, value in preserved.items(): + if value is None: + continue + # The _transform_for_* helpers never set these clauses on `new`, so + # re-attaching a preserved clause can never clobber one `new` supplies. + # Assert that invariant rather than silently choosing one — a future + # rewrite that emits its own copy must reconcile the two explicitly (for + # WITH, by merging the CTE lists rather than dropping either). + assert new.args.get(key) is None, ( + f"transplant() cannot reconcile a {key!r} clause from `new` with the " + "preserved outer one; a rewrite that emits its own must merge them." + ) + select.set(key, value) diff --git a/src/giql/expanders/_params.py b/src/giql/expanders/_params.py new file mode 100644 index 0000000..60eaa83 --- /dev/null +++ b/src/giql/expanders/_params.py @@ -0,0 +1,37 @@ +"""Shared operator-parameter coercions for the expanders. + +The private ``_``-prefix keeps this module out of the ``expanders`` package +auto-discovery (see :mod:`giql.expanders`), so it is a plain helper module rather +than a registered expander. +""" + +from __future__ import annotations + +from sqlglot import exp + + +def coerce_bool_param(param: exp.Expression | None) -> bool: + """Coerce an optional operator boolean argument to a Python ``bool``. + + The ``stranded`` / ``signed`` keyword arguments of DISTANCE and NEAREST are + read the same way: an absent argument is ``False``, an + :class:`sqlglot.exp.Boolean` node yields ``bool(node.this)``, and anything + else is truthy when its rendered text is ``TRUE`` / ``1`` / ``YES``. The + :func:`bool` coercion guarantees a real ``bool`` regardless of what the + parser stored on the node. + + Parameters + ---------- + param : exp.Expression | None + The ``stranded`` or ``signed`` argument node, or ``None`` if absent. + + Returns + ------- + bool + The coerced Python boolean (``False`` when the argument is absent). + """ + if not param: + return False + if isinstance(param, exp.Boolean): + return bool(param.this) + return str(param).upper() in ("TRUE", "1", "YES") diff --git a/src/giql/expanders/_per_chrom.py b/src/giql/expanders/_per_chrom.py new file mode 100644 index 0000000..be5f96b --- /dev/null +++ b/src/giql/expanders/_per_chrom.py @@ -0,0 +1,74 @@ +"""Shared DuckDB per-chromosome dynamic-SQL scaffolding (#217). + +Several DuckDB operator overrides reach the fast ``IE_JOIN`` range-join operator +by partitioning a range join per chromosome, which removes the low-cardinality +``chrom`` equality that would otherwise force a ``HASH_JOIN``. Because GIQL does +not know a table's chromosomes at transpile time, the per-chromosome +``UNION ALL`` is assembled at runtime: a ``string_agg`` over the distinct +chromosomes builds the dynamic SQL into a session variable, which the outer query +reads back through ``query(getvariable(...))``. + +This module centralizes the target-agnostic scaffolding — the chromosome +string-literal interpolation, a collision-resistant variable name, the +``SET VARIABLE`` builder, and the dynamic-relation reference. The +operator-specific parts (the per-chromosome SQL template, the chromosome +partition source, the empty-schema fallback, and the outer projection) stay with +each expander. First extracted from +:mod:`giql.expanders.intersects_duckdb`; reused by DISJOIN (#216) and NEAREST. +""" + +from __future__ import annotations + +from uuid import uuid4 + +#: A DuckDB expression that renders the ``string_agg`` partition's ``chrom`` +#: column as a single-quoted, injection-safe SQL string literal, for +#: interpolation into a per-chromosome ``WHERE chrom = ''`` filter. +CHROM_LITERAL = "'''' || replace(chrom, '''', '''''') || ''''" + + +def sql_escape(text: str) -> str: + """Return *text* with single quotes doubled for safe SQL string-literal use.""" + return text.replace("'", "''") + + +def new_variable_name(prefix: str) -> str: + """Return a collision-resistant DuckDB session-variable name. + + The full uuid4 hex (128 bits) keeps the name unique even across many + ``transpile()`` calls interleaved in one DuckDB session -- session variables + are global session state, so a token collision would silently rebind a + wrapper query to a different dynamic relation. + """ + return f"{prefix}_{uuid4().hex}" + + +def set_variable_statement( + var_name: str, + per_chrom_template: str, + chrom_partition: str, + empty_schema: str, +) -> str: + """Build the ``SET VARIABLE`` statement that assembles the per-chrom dynamic SQL. + + *per_chrom_template* is a SQL string-concat expression that ``string_agg`` + evaluates once per row of *chrom_partition* (a ``SELECT DISTINCT `` + source), producing one per-chromosome ``SELECT`` joined by ``UNION ALL``. + When no chromosome passes the partition the aggregate is ``NULL``, so it + falls back to *empty_schema* -- a ``... WHERE FALSE`` query that lets DuckDB + resolve every output column to its real declared type. + """ + return ( + f"SET VARIABLE {var_name} = COALESCE((\n" + f" SELECT string_agg(\n" + f" {per_chrom_template},\n" + f" ' UNION ALL '\n" + f" )\n" + f" FROM ({chrom_partition})\n" + f"), '{sql_escape(empty_schema)}')" + ) + + +def dynamic_relation(var_name: str, alias: str) -> str: + """Return the ``query(getvariable('')) AS `` relation reference.""" + return f"query(getvariable('{var_name}')) AS {alias}" diff --git a/src/giql/expanders/cluster.py b/src/giql/expanders/cluster.py new file mode 100644 index 0000000..42e947e --- /dev/null +++ b/src/giql/expanders/cluster.py @@ -0,0 +1,644 @@ +"""The CLUSTER operator expander (epic #137, issue #144). + +CLUSTER assigns a cluster id to every interval, grouping each run of mutually +adjacent intervals (optionally within a maximum gap, strand-aware, and gated by a +pairwise predicate) under one id. It cannot be a single window function because it +needs a window over a window (``SUM`` over a running-edge-derived boundary flag), +so the expansion restructures the enclosing query into a two-level form:: + + SELECT *, CLUSTER(interval) AS cluster_id FROM features + +becomes:: + + SELECT *, + SUM(__giql_is_new_cluster) OVER (PARTITION BY chrom ORDER BY start) + AS cluster_id + FROM ( + SELECT *, + CASE WHEN MAX(end) OVER (...) >= start THEN 0 ELSE 1 END + AS __giql_is_new_cluster + FROM features + ) AS __giql_lag_calc + +The boundary flag keys off the running maximum end of the preceding rows (the +cluster's right edge so far), not the immediately preceding row's end, so a later +interval contained within an earlier wider one is not spuriously split (#214). + +(The ``becomes::`` form is simplified for readability; the emitted SQL quotes +identifiers, appends ``NULLS LAST`` to the window ORDER BY and a ``ROWS BETWEEN +UNBOUNDED PRECEDING AND 1 PRECEDING`` frame, and — for a bare or qualified star +projection — emits the outer star as ``* EXCEPT (__giql_is_new_cluster)`` to hide +the synthesized flag from the output (#184, #185).) + +This module is the AST-expansion replacement for the legacy +``ClusterTransformer``, which ran as a pre-pass transformer on the raw parsed AST +(in the former ``giql.transformer`` module, removed once every join rewrite became +registry-driven — #144, #169). It produces the same SQL; the existing CLUSTER +transpilation and bedtools oracle tests are the migration oracle. + +Unlike the node-local expanders for DISTANCE / NEAREST / DISJOIN — whose result +replaces a single node — CLUSTER is a **whole-query rewrite**. It shares only the +*shape* of :func:`giql.expanders.nearest._fallback_form` — return the operator +node unchanged so the pass's ``node.replace`` is a no-op — but the *mechanism* +differs. ``nearest`` does an in-place child ``.replace`` because its LATERAL has a +parent to replace through; CLUSTER instead copies the enclosing +:class:`~sqlglot.expressions.Select`, restructures the copy, and *transplants* its +contents back onto the original SELECT (see +:func:`giql.expanders._genomic.transplant`). That root- +preserving transplant is required because the canonical +``SELECT *, CLUSTER(...) FROM t`` puts CLUSTER at the *root* ``SELECT``, which has +no parent to replace it through; transplanting preserves the root's identity (the +pass returns the same expression object it was handed). + +The pass walks the tree and collects every operator node, expanding deepest-first, +so this expander never recurses into CTEs / subqueries: a nested CLUSTER is +collected and expanded on its own. The shared cluster-restructure helper +(:func:`expand_cluster_query`) is reused by :mod:`giql.expanders.merge`, since +MERGE is built on CLUSTER. +""" + +from __future__ import annotations + +from sqlglot import exp + +from giql.constants import CLUSTER_FLAG_COL +from giql.constants import CLUSTER_SIBLING_PREFIX +from giql.expander import ExpansionContext +from giql.expander import expand_operators +from giql.expander import register +from giql.expanders._genomic import GenomicColumns +from giql.expanders._genomic import extract_stranded +from giql.expanders._genomic import find_projected +from giql.expanders._genomic import genomic_columns +from giql.expanders._genomic import is_star_projection +from giql.expanders._genomic import reject_cluster_merge_mix +from giql.expanders._genomic import require_top_level_projection +from giql.expanders._genomic import transplant +from giql.expressions import GIQLCluster +from giql.targets import GenericTarget + + +@register(GenericTarget, GIQLCluster) +def expand_cluster(node: GIQLCluster, ctx: ExpansionContext) -> exp.Expression: + """Expand a CLUSTER node by restructuring its enclosing SELECT in place. + + Registered for :class:`~giql.targets.GenericTarget`, so every target resolves + to it through the registry's generic chain (CLUSTER emits identical SQL across + targets). Locates the enclosing :class:`~sqlglot.expressions.Select`, derives + the genomic columns from the FROM table, and rewrites that SELECT in place into + the two-level ``__giql_lag_calc`` form. Returns the original CLUSTER node (now + unreachable from the rewritten root SELECT) so the pass's ``node.replace`` is a + no-op. + + Parameters + ---------- + node : GIQLCluster + The CLUSTER node being expanded. + ctx : ExpansionContext + The expansion context; only its :attr:`~ExpansionContext.tables` is read + (CLUSTER derives its columns from the enclosing FROM table, not from + resolution metadata). + + Returns + ------- + exp.Expression + The CLUSTER node, unchanged — the surrounding SELECT was mutated in place. + """ + select = node.find_ancestor(exp.Select) + if select is None: + # Defensive: CLUSTER only ever parses inside a SELECT projection. With no + # enclosing SELECT there is nothing to restructure; leave the node for the + # generator to error on, exactly as the legacy transformer's + # non-Select guard did. + return node + require_top_level_projection(select, node, GIQLCluster) + reject_cluster_merge_mix(select) + columns = genomic_columns(select, ctx.tables) + # Build the transformed query from a detached copy so the intermediate never + # aliases live nodes, then transplant its args onto the original SELECT to + # preserve its identity (and so a root SELECT is rewritten without a parent to + # replace it through). + source = select.copy() + # Hiding the synthesized __giql_is_new_cluster flag from a surfacing star is the + # safe default (#184; MERGE opts out — see expand_cluster_query). This runs for + # every CLUSTER the pass expands, root or nested: harmless when the flag would + # not surface, and it also stops a nested CLUSTER's flag leaking through an + # enclosing star. + transformed = expand_cluster_query(source, columns) + if transformed is None: + # No-op for a CLUSTER that parses outside the SELECT projection (e.g. in + # WHERE / ORDER BY): find_projected finds none in the copy, so there is + # nothing to restructure and the operator leaks to the generator exactly + # as on `main`. require_top_level_projection above already rejects the + # in-projection-expression case; this guards the out-of-projection case. + return node + transplant(select, transformed) + # copy()+transplant duplicated the enclosing WHERE / HAVING into the new + # __giql_lag_calc subquery; the originals the pass collected are now unreachable, so + # re-run the pass over the restructured SELECT to expand any sibling pass-3 + # operators (spatial predicates, DISTANCE) carried into it. Safe from + # recursion: the CLUSTER node is already replaced by its SUM window. (#144 B1) + # + # expand_operators may return a new root (a registered statement finalizer can + # wrap it), and its contract requires callers to use the return value rather + # than assume in-place mutation, so reinstall a new root in place of `select`. + # Today the branch is never taken here — the sole finalizer-registering operator + # (a correlated NEAREST fallback) is already a plain join by this deepest-first + # re-walk, so this nested run registers no finalizer and `result is select` in + # practice — but honoring the contract keeps the seam future-proof. A NEAREST + # fallback nested under this CLUSTER registered its finalizer in the *outer* + # expand_operators run; that finalizer re-locates the transplanted join by its + # reserved `meta` tag and wraps its enclosing SELECT after this rewrite, so a + # `SELECT *`/`b.*` that would re-surface the fallback's reserved columns is + # projected away rather than leaking (#172). + result = expand_operators(select, ctx.target, ctx.tables, ctx.registry) + if result is not select: + select.replace(result) + return node + + +def expand_cluster_query( + query: exp.Select, columns: GenomicColumns, hide_reserved: bool = True +) -> exp.Select | None: + """Restructure *query* for every CLUSTER in its projection; ``None`` if none. + + Finds CLUSTER expressions in *query*'s SELECT list and rewrites the query into + the two-level ``__giql_lag_calc`` form once per CLUSTER (chaining, as the legacy + transformer did). Operates on *query* directly and returns the rewritten + query. Returns ``None`` when *query* projects no CLUSTER, so callers can treat + that as a no-op. + + *hide_reserved* (default ``True``) rewrites a star in the outer projection to hide + the synthesized ``__giql_is_new_cluster`` flag the star would otherwise surface + (#184). Hiding is the default so a caller cannot silently reintroduce the leak by + omitting it. MERGE — which reuses this helper to build its intermediate clustered + query — is the one caller that opts out (``hide_reserved=False``): its explicit + outer projection never surfaces the flag, so hiding it would only add a needless + ``* EXCEPT`` to the intermediate ``__giql_clustered`` subquery. + """ + cluster_exprs = find_projected(query, GIQLCluster) + if not cluster_exprs: + return None + if len(cluster_exprs) > 1: + # Chaining the rewrite per CLUSTER yields a duplicate ``__giql_lag_calc`` + # alias and an ``__giql_is_new_cluster`` binder error — non-executable SQL. + # Fail loudly, mirroring the multiple-MERGE guard, rather than emitting the + # broken SQL (#144 A15). + raise ValueError("Multiple CLUSTER expressions not yet supported") + if sum(is_star_projection(item) for item in query.expressions) > 1: + # The two-level rewrite copies the whole projection into the inner + # ``__giql_lag_calc`` subquery and re-expands each star at the outer level, so + # N stars multiply the base columns (N inner expansions × N outer), silently + # producing the wrong column set. A single star is the supported shape; fail + # loudly on multiple, mirroring the multiple-CLUSTER guard above, rather than + # emitting a silently multiplied projection (#194). (MERGE reuses this helper + # with a synthesized single-star intermediate, so this never fires there.) + raise ValueError( + "CLUSTER does not support multiple star projections " + "(e.g. SELECT *, *, CLUSTER(...)); project a single star" + ) + for cluster_expr in cluster_exprs: + query = _transform_for_cluster(query, cluster_expr, columns, hide_reserved) + return query + + +def _transform_for_cluster( + query: exp.Select, + cluster_expr: GIQLCluster, + columns: GenomicColumns, + hide_reserved: bool = True, +) -> exp.Select: + """Rewrite *query* into the two-level ``__giql_lag_calc`` form for one CLUSTER. + + A byte-for-byte port of the legacy + ``ClusterTransformer._transform_for_cluster``, with the genomic columns passed + in (the legacy method re-derived them from the FROM table) rather than read off + ``self``. Builds an inner ``__giql_lag_calc`` subquery that materializes an + ``__giql_is_new_cluster`` flag from a ``LAG`` window, then an outer query whose + ``SUM(__giql_is_new_cluster) OVER (...)`` window replaces the CLUSTER projection. + """ + chrom_col, start_col, end_col, strand_col = columns + + # Extract CLUSTER parameters + distance_expr = cluster_expr.args.get("distance") + + # Handle distance parameter - could be int literal or None + if distance_expr: + if isinstance(distance_expr, exp.Literal): + distance = int(distance_expr.this) + else: + # Defensive: the grammar only yields exp.Literal for a distance + # argument, so this non-Literal fallback is unreachable via the parser; + # retained for parity with the legacy port. + try: + distance = int(str(distance_expr.this)) + except (ValueError, AttributeError): + distance = 0 + else: + distance = 0 + + stranded = extract_stranded(cluster_expr.args.get("stranded")) + + # Build partition clause + partition_cols = [exp.column(chrom_col, quoted=True)] + if stranded: + partition_cols.append(exp.column(strand_col, quoted=True)) + + # Build ORDER BY for window + order_by = [exp.Ordered(this=exp.column(start_col, quoted=True))] + + # Cluster boundaries key off the running maximum end of the preceding rows -- + # the cluster's right edge so far -- NOT the immediately preceding row's end. + # LAG(end) would spuriously split a later interval that still overlaps a + # cluster whose previous row is a contained interval ending early (#214). + running_max_end = exp.Window( + this=exp.Max(this=exp.column(end_col, quoted=True)), + partition_by=partition_cols, + order=exp.Order(expressions=order_by), + spec=exp.WindowSpec( + kind="ROWS", + start="UNBOUNDED", + start_side="PRECEDING", + end="1", + end_side="PRECEDING", + ), + ) + + # Add distance offset if specified + if distance > 0: + edge_with_distance = exp.Add( + this=running_max_end, expression=exp.Literal.number(distance) + ) + else: + edge_with_distance = running_max_end + + # Build the adjacency condition (running cluster edge >= current start). + adjacency = exp.GTE( + this=edge_with_distance, + expression=exp.column(start_col, quoted=True), + ) + + # An optional predicate further restricts which adjacent intervals are kept + # together: a row stays in the current cluster only when it is adjacent to the + # cluster's running-max edge AND the predicate holds against its immediate + # sorted predecessor. ``PREV(col)`` references in the predicate resolve to that + # predecessor row via LAG over the same partition/order. Note the two notions + # can reference different rows under containment -- adjacency comes from an + # earlier, wider interval while the predicate compares to the immediate + # predecessor -- matching the documented immediate-predecessor semantics (#214). + predicate_expr = cluster_expr.args.get("predicate") + if predicate_expr is not None: + rewritten_predicate = _rewrite_predecessor_refs( + predicate_expr, partition_cols, order_by + ) + keep_together = exp.And( + this=adjacency, + expression=exp.Paren(this=rewritten_predicate), + ) + else: + keep_together = adjacency + + # Create CASE expression for __giql_is_new_cluster + case_expr = exp.Case( + ifs=[ + exp.If( + this=keep_together, + true=exp.Literal.number(0), + ) + ], + default=exp.Literal.number(1), + ) + + # When the projection carries a star, every sibling item (each non-star, non-CLUSTER + # projection alongside it) is re-projected at the outer level by *reference* to the + # column the inner subquery materializes (see the outer loop below), pinned to its + # written position. Referencing it by the user's own name is unsafe: the name can + # collide with a base column the star also surfaces (``SELECT *, expr AS score`` — + # the outer ``EXCEPT`` would strip the base ``score`` and the reference would bind + # ambiguously) or with another sibling's name (``SELECT *, 1 AS x, 2 AS x`` — a + # duplicate ``EXCEPT`` entry, a hard parse error). Materialize each sibling in the + # inner subquery under a reserved, unique ``__giql_sibling_N`` name instead; the + # outer star EXCEPTs those reserved names and the outer projection references them, + # aliased back to the name they would carry un-clustered, at their written position — + # so base columns are never stripped, the inner window references stay unambiguous, + # and both the values *and the column order* match the un-clustered projection. + # Materializing *every* sibling — not only aliased ones — is what pins column order: + # a non-aliased sibling left to ride the star would always surface at the star's + # position, silently reordering it ahead of any aliased sibling written before it + # (#190; column-order pinning #192 round 3). + outer_has_star = any(is_star_projection(item) for item in query.expressions) + sibling_inner_names: dict[int, str] = {} + if outer_has_star: + for item in query.expressions: + if is_star_projection(item): + continue + if isinstance(item, GIQLCluster): + continue + if isinstance(item, exp.Alias) and isinstance(item.this, GIQLCluster): + continue + sibling_inner_names[id(item)] = ( + f"{CLUSTER_SIBLING_PREFIX}{len(sibling_inner_names)}" + ) + + # Build CTE SELECT expressions (all original except CLUSTER, plus + # __giql_is_new_cluster) + cte_expressions = [] + for expression in query.expressions: + # Skip CLUSTER expressions + if isinstance(expression, GIQLCluster): + continue + elif isinstance(expression, exp.Alias) and isinstance( + expression.this, GIQLCluster + ): + continue + elif id(expression) in sibling_inner_names: + # A star's sibling: materialize it under its reserved inner name so the outer + # star can EXCEPT it (without stripping a same-named base column) and the + # outer projection can re-project it at the user's written slot. Materialize + # the aliased item's inner value, or the bare column / unnamed expression + # itself when it has no alias. + value = expression.this if isinstance(expression, exp.Alias) else expression + cte_expressions.append( + exp.alias_( + value.copy(), + sibling_inner_names[id(expression)], + quoted=False, + ) + ) + else: + cte_expressions.append(expression) + + # Ensure required columns for window functions are included + required_cols = {chrom_col, start_col, end_col} + if stranded: + required_cols.add(strand_col) + + # The predicate is evaluated inside the __giql_lag_calc CTE, so every column + # it references (current-row columns and PREV() arguments alike) must + # be projected into that CTE. Folding them into required_cols makes the + # scope dependency explicit and keeps the columns available even when a + # later operator wraps this query in a further subquery. + if predicate_expr is not None: + required_cols |= {col.name for col in predicate_expr.find_all(exp.Column)} + + # Check if required columns are already in the select list + selected_cols = set() + for expr in cte_expressions: + if is_star_projection(expr): + # A bare ``*`` or a qualified ``rel.*`` covers every column. A qualified + # ``rel.*`` parses as a Column wrapping a Star, so it must be caught here, + # before the plain-Column branch below — whose ``.name`` for a star is the + # literal ``'*'`` (never a required-column name), so it would leave chrom/ + # start/end seen as missing and re-inject them as chrom_1/start_1/end_1 + # duplicates (#185). + selected_cols = required_cols # Assume all are covered + break + elif isinstance(expr, exp.Column): + selected_cols.add(expr.name) + elif isinstance(expr, exp.Alias): + # Don't count aliases as the source column + pass + + # Add missing required columns + # Sort the residual so the injected-column order is deterministic across runs + # (set-difference iteration order is PYTHONHASHSEED-dependent; the legacy port + # left it nondeterministic — see review #144 A2). + for col in sorted(required_cols - selected_cols): + cte_expressions.append(exp.column(col, quoted=True)) + + # Add __giql_is_new_cluster calculation. The reserved __giql_ prefix keeps the + # synthesized flag from colliding with a user column of the same name (#161). + cte_expressions.append(exp.alias_(case_expr, CLUSTER_FLAG_COL, quoted=False)) + + # Build CTE query + cte_select = exp.Select() + cte_select.select(*cte_expressions, copy=False) + + # Copy FROM, WHERE, GROUP BY, HAVING from original (but not ORDER BY) + # Use copy() to avoid sharing references between queries + if query.args.get("from_"): + from_clause = query.args["from_"].copy() + cte_select.set("from_", from_clause) + if query.args.get("where"): + cte_select.set("where", query.args["where"].copy()) + if query.args.get("group"): + cte_select.set("group", query.args["group"].copy()) + if query.args.get("having"): + cte_select.set("having", query.args["having"].copy()) + + # Create outer query with SUM over __giql_is_new_cluster + sum_window = exp.Window( + this=exp.Sum(this=exp.column(CLUSTER_FLAG_COL)), + partition_by=partition_cols, + order=exp.Order(expressions=order_by), + ) + + # When the projection carries a star, every sibling was materialized in the inner + # __giql_lag_calc subquery under a reserved ``__giql_sibling_N`` name (the CTE loop + # above). The outer star EXCEPTs those reserved names so it does not surface them, + # and each sibling is re-projected at its written position as a reference to its + # reserved column, aliased back to the output name it would carry un-clustered. + # Referencing the materialized column rather than re-evaluating the original + # expression keeps it correct even when that expression's inputs were EXCEPTed from + # the star inside the subquery, and — because the reserved name cannot collide — + # keeps it correct for any name. Re-projecting *every* sibling (not only aliased + # ones) pins each to the user's column order (#190; column-order pinning #192). + # + # Without a star nothing re-surfaces the item, so it is always kept as-is. + star_hidden_names = list(sibling_inner_names.values()) + + # Build outer SELECT expressions (replace CLUSTER with SUM) + new_expressions = [] + for expression in query.expressions: + if isinstance(expression, GIQLCluster): + new_expressions.append(sum_window) + elif isinstance(expression, exp.Alias) and isinstance( + expression.this, GIQLCluster + ): + # Keep the alias but replace the expression + new_expressions.append( + exp.alias_(sum_window, expression.alias, quoted=False) + ) + elif is_star_projection(expression): + # A star projection (bare ``*`` or qualified ``rel.*``) expands over the + # inner __giql_lag_calc subquery. It is normalized to a *bare* outer star + # (dropping any dangling ``rel.`` qualifier — an executability fix), the + # synthesized __giql_is_new_cluster flag is EXCEPTed from it when + # hide_reserved is set (the default; MERGE opts out — see + # expand_cluster_query), and each aliased sibling re-projected below is + # EXCEPTed too so the star does not double-surface it (#190). See + # _outer_star_over_lag_calc for how the concerns are gated. + rewritten = _outer_star_over_lag_calc( + expression, CLUSTER_FLAG_COL, hide_reserved, star_hidden_names + ) + new_expressions.append(expression if rewritten is None else rewritten) + elif not outer_has_star: + # A non-star, non-CLUSTER item with no sibling star: keep it unchanged. + new_expressions.append(expression) + elif id(expression) in sibling_inner_names: + # A sibling of a star: re-project the reserved column the CTE materialized it + # under, aliased back to the output name it would carry un-clustered, at its + # written position — so it surfaces exactly once (its reserved name is + # EXCEPTed from the star above), collision-free for any name, in the + # user's column order rather than jumping to the star's position. The output + # name is the user's alias, a bare column's own name, or — for an unnamed + # expression, which carries no output name — its rendered text, so a + # positional consumer sees it in the written slot (#190; ordering #192). + out_name = expression.output_name or expression.sql() + new_expressions.append( + exp.alias_( + exp.column(sibling_inner_names[id(expression)], quoted=True), + out_name, + quoted=True, + ) + ) + + # Build new query + new_query = exp.Select() + new_query.select(*new_expressions, copy=False) + + # Wrap CTE in subquery and set as FROM clause. The alias carries the reserved + # __giql_ prefix for namespace consistency with the other synthesized names + # (#161); a derived-table alias never actually collides with a user relation + # (it lives in the enclosing query's scope), so this is hygiene, not a fix. + subquery = exp.Subquery( + this=cte_select, + alias=exp.TableAlias(this=exp.Identifier(this="__giql_lag_calc")), + ) + new_query.from_(subquery, copy=False) + + # Copy ORDER BY from original to outer query + if query.args.get("order"): + new_query.order_by(*query.args["order"].expressions, copy=False) + + return new_query + + +def _outer_star_over_lag_calc( + expression: exp.Expression, + reserved: str, + hide_reserved: bool, + extra_hidden: list[str] | None = None, +) -> exp.Star | None: + """Rewrite a star projection for the outer query over ``__giql_lag_calc``. + + Returns a *bare* :class:`~sqlglot.expressions.Star` for a bare ``*`` or a qualified + ``rel.*`` outer projection, or ``None`` for any non-star item. (The caller only ever + invokes this on the star branch, so the ``None`` return is defensive — non-star + items are handled by the caller, not passed here.) Three concerns are handled, on + separate gates: + + **Normalization to a bare star — always.** The CLUSTER rewrite's outer query selects + from the single renamed ``__giql_lag_calc`` subquery, so the user's ``rel.`` + qualifier no longer resolves there — an un-rewritten ``rel.*`` would dangle + (``Binder Error: Referenced table "rel" not found``). CLUSTER runs over a single + relation (joins are rejected upstream by + :func:`giql.expanders._genomic.genomic_columns`), so every column ``rel.*`` would + have named is exactly the subquery's projection, making ``rel.*`` equivalent to + ``*``. The qualifier is therefore *always* dropped to a bare star. This is an + **executability** fix (without it the query does not bind), so it must NOT be gated + on the cosmetic *hide_reserved* flag (#185). Any ``EXCEPT`` the user's own star + carried is already applied *inside* ``__giql_lag_calc`` (the inner subquery keeps + the original projection node), so the outer bare star never needs to carry it. + + **Flag-hiding — only when *hide_reserved*.** The subquery's ``*`` also carries the + synthesized ``reserved`` flag (``__giql_is_new_cluster``), so a bare outer ``*`` + would surface it — and, under a preserved DISTINCT, dedup over it. When + *hide_reserved* is set (the default; MERGE opts out — its explicit projection never + surfaces the flag), the bare star is emitted as ``* EXCEPT (reserved)`` to hide it + (#184). + + The hiding is applied *eagerly*, at the same projection level as the outer star — + NOT via an outer-wrapping statement finalizer like NEAREST's + (:func:`giql.expanders.nearest._make_reserved_column_finalizer`). It must be: + ``transplant`` re-attaches the user's ``DISTINCT`` onto this outer query, so a + finalizer that wrapped it in ``SELECT * EXCEPT (...) FROM ()`` would dedup + over the flag *before* stripping it (splitting rows that should collapse). NEAREST + can wrap because its reserved columns surface through an arbitrary enclosing SELECT + it does not own; CLUSTER owns its outer star at build time and must excise the flag + in place. A future consolidation (#187) must preserve this level constraint. + + **Aliased-sibling hiding — when *extra_hidden* is given.** When the projection pairs + the star with aliased sibling items (``SELECT *, expr AS name, CLUSTER(...)``), each + sibling was materialized in ``__giql_lag_calc`` under a reserved ``__giql_sibling_N`` + name and would be surfaced a second time by the outer star. The caller re-projects + each such sibling (aliased back to the user's name) at its own position and passes + the reserved names as *extra_hidden* so they are EXCEPTed from this star too — the + star then carries only the base columns, and the sibling surfaces exactly once, in + the user's column order. The reserved names are used precisely so this EXCEPT can + never strip a same-named base column the star should keep (#190). + + ``* EXCEPT`` is the portable exclusion form GIQL already relies on (the coordinate + canonicalizer and the NEAREST fallback finalizer emit it too), so sqlglot renders + the target spelling (``EXCLUDE`` for DuckDB) without a capability branch here. + """ + if not is_star_projection(expression): + return None + hidden = [] + if hide_reserved: + hidden.append(exp.column(reserved, quoted=True)) + if extra_hidden: + hidden.extend(exp.column(name, quoted=True) for name in extra_hidden) + if hidden: + return exp.Star(except_=hidden) + return exp.Star() + + +def _rewrite_predecessor_refs( + predicate: exp.Expression, + partition_cols: list[exp.Expression], + order_by: list[exp.Ordered], +) -> exp.Expression: + """Rewrite ``PREV(col)`` calls in a predicate to LAG windows. + + A byte-for-byte port of the legacy + ``ClusterTransformer._rewrite_predecessor_refs``. Bare column references in the + predicate denote the current interval. Each ``PREV(col)`` call denotes the + sorted predecessor's value of that column and is rewritten to ``LAG(col) OVER + (...)`` using the same partition/order as the cluster's adjacency window, so + the predicate is evaluated pairwise against the immediately preceding row. + Every column identifier (current-row columns and LAG arguments alike) is quoted + so that reserved-word genomic columns such as ``start`` / ``end`` are emitted as + valid SQL. + + :param predicate: + Boolean predicate expression to rewrite (not mutated). + :param partition_cols: + Window partition columns (chromosome, optionally strand). + :param order_by: + Window ORDER BY terms (start position). + :return: + A copy of the predicate with every ``PREV(...)`` call replaced by an + equivalent LAG window and all column identifiers quoted. + :raises ValueError: + If a ``PREV()`` call does not take exactly one argument, or if a + ``PREV()`` call is nested inside another (predicates compare only the + immediate predecessor). + """ + + def _is_prev(node: exp.Expression) -> bool: + return isinstance(node, exp.Anonymous) and node.name.upper() == "PREV" + + def _replace(node: exp.Expression) -> exp.Expression: + if _is_prev(node): + args = node.expressions + if len(args) != 1: + raise ValueError( + f"PREV() takes exactly one column argument; got {len(args)}." + ) + if any(_is_prev(inner) for inner in args[0].find_all(exp.Anonymous)): + raise ValueError( + "PREV() cannot be nested; a CLUSTER/MERGE predicate " + "compares only the immediate predecessor." + ) + return exp.Window( + this=exp.Anonymous(this="LAG", expressions=[args[0].copy()]), + partition_by=[col.copy() for col in partition_cols], + order=exp.Order(expressions=[term.copy() for term in order_by]), + ) + return node + + rewritten = predicate.copy().transform(_replace) + for column in rewritten.find_all(exp.Column): + column.this.set("quoted", True) + return rewritten diff --git a/src/giql/expanders/disjoin.py b/src/giql/expanders/disjoin.py new file mode 100644 index 0000000..f0c178c --- /dev/null +++ b/src/giql/expanders/disjoin.py @@ -0,0 +1,348 @@ +"""The DISJOIN operator expander (epic #137, wave 3 — issue #143). + +DISJOIN splits each target interval at every reference breakpoint strictly +interior to it, passing the full target row through and appending the +sub-interval as ``disjoin_chrom`` / ``disjoin_start`` / ``disjoin_end`` in the +target table's declared coordinate system. A coverage filter drops sub-intervals +overlapping no reference interval; an omitted reference defaults to the target +set. + +This module is the AST-expansion replacement for the legacy +``giqldisjoin_sql`` generator emitter. It assembles the same WITH-CTE +subquery, parses it back into a sqlglot :class:`~sqlglot.expressions.Expression`, +and returns that node so the active target's serializer renders it — dissolving +the emit-time string special-case left by epic #114. + +Three behaviours are capability-driven / corrected here relative to the legacy +emitter: + +* **star-REPLACE portability (#143).** The full-row passthrough de-canonicalizes + the interval columns back into the target's declared encoding. On a target that + supports ``SELECT * REPLACE (...)`` (``supports_star_replace`` — DuckDB) the + passthrough emits ``t.* REPLACE (...)``. On a target without it (the generic + baseline and the DataFusion family) it emits the portable + ``t.* EXCEPT (...)`` form plus the two recomputed interval columns. This + portable form is **not** SQL-92 — it requires ``SELECT * EXCEPT (...)`` + support, which the DataFusion family provides but a strict SQL-92 engine (and + DuckDB, which prefers ``* REPLACE``) does not — so the generic non-canonical + passthrough is **not DuckDB-runnable** and runs only on an ``* EXCEPT``-capable + engine. +* **passthrough column order diverges across targets (issue #143).** The two + passthrough forms are *row-equivalent* but **not column-order-equivalent**. + ``* REPLACE`` substitutes the start/end columns *in place*, so the row keeps + its original column order. ``* EXCEPT (start, end)`` drops those two columns + and the recomputed ones are *re-appended* at the end of the projection, so on + the generic / DataFusion path the start/end columns move to the right of any + passthrough column. A ``SELECT *`` over a non-canonical DISJOIN is therefore + column-order-divergent between DuckDB and the generic target; only the row + *set* is identical. Select **explicit** columns (in a chosen order) when the + output column order must agree across targets. +* **duplicate-column fix (#153).** Every UNION branch of the ``__giql_dj_cuts`` + CTE aliases all four projected columns (``kc`` / ``ks`` / ``ke`` / ``pos``), so + the ``t."end"`` column and the end-cut expression never collide on an output + name when ``end`` de-canonicalizes to the bare physical column (the default + 0-based half-open identity case). DuckDB tolerated the prior unaliased + duplicate; DataFusion rejected it as a non-unique projection name. + +Input canonicalization is owned by ``CanonicalizeCoordinates`` (pass 2, #122): +every non-canonical interval-bearing operand is rewritten to a canonical +``__giql_canon_*`` CTE before this pass runs, so the expander consumes +already-canonical 0-based half-open columns and applies no input-canonicalization +arithmetic. The output round-trip back to the target's declared encoding stays +here, driven by the original encoding pass 2 preserves on the resolution. +""" + +from __future__ import annotations + +from sqlglot import exp +from sqlglot import parse_one + +from giql.canonical import decanonical_end +from giql.canonical import decanonical_start +from giql.constants import DJ_PREFIX +from giql.dialect import GIQLDialect +from giql.expander import ExpansionContext +from giql.expander import register +from giql.expressions import GIQLDisjoin +from giql.resolver import OperatorResolution +from giql.resolver import ResolvedRef +from giql.table import Table +from giql.targets import GenericTarget + + +@register(GenericTarget, GIQLDisjoin) +def expand_disjoin(node: GIQLDisjoin, ctx: ExpansionContext) -> exp.Expression: + """Expand a DISJOIN node to its WITH-CTE subquery AST for the active target. + + Registered for :class:`~giql.targets.GenericTarget`, so it is the portable + fallback every target resolves to through the registry's generic chain. The + one capability branch — ``ctx.capabilities.supports_star_replace`` — selects + the passthrough projection form, so a single expander covers both the + ``* REPLACE`` (DuckDB) and ``* EXCEPT`` (DataFusion / generic) targets. + + Parameters + ---------- + node : GIQLDisjoin + The :class:`~giql.expressions.GIQLDisjoin` node being expanded. + ctx : ExpansionContext + The expansion context carrying the node's pass-1/pass-2 resolution + metadata, the active target and its capabilities, and the registered + tables. + + Returns + ------- + exp.Expression + The parsed AST of the parenthesized WITH-CTE subquery that replaces the + DISJOIN node. + """ + sql = _build_disjoin_sql(node, ctx) + return parse_one(sql, dialect=GIQLDialect) + + +def _build_disjoin_sql(node: GIQLDisjoin, ctx: ExpansionContext) -> str: + """Assemble the DISJOIN WITH-CTE subquery SQL string for the active target. + + Mirrors the legacy ``giqldisjoin_sql`` emitter's shape one named + ``__giql_dj_*`` CTE at a time, with the #153 alias fix and the + capability-driven passthrough applied. + """ + resolution = ctx.resolution + target_ref, ref, ref_from = _disjoin_resolution(node, resolution) + target_name = target_ref.name + target_chrom, target_start, target_end = target_ref.cols + + ref_chrom, ref_start, ref_end = ref.cols + is_self_reference = ref.coverage_skippable + + # The target's *declared* encoding, which disjoin_* output and the + # passed-through interval round-trip back into. Pass 2 preserves it on the + # resolution when it wraps a non-canonical target (the slot's own Table is + # then None); a canonical target keeps the (identity) encoding on its slot. + output_table = _disjoin_output_encoding(resolution, target_ref) + + # Post-pass every operand is canonical 0-based half-open, so the physical + # columns are consumed verbatim with no input-canonicalization arithmetic. + t_chrom = f't."{target_chrom}"' + t_start = f't."{target_start}"' + t_end = f't."{target_end}"' + + # Reference endpoints: unqualified for the breakpoint CTE, qualified by 'r' + # for the coverage EXISTS filter. + bp_start = f'"{ref_start}"' + bp_end = f'"{ref_end}"' + r_start = f'r."{ref_start}"' + r_end = f'r."{ref_end}"' + + # disjoin_start / disjoin_end are emitted in the target's declared coordinate + # system so an output row carries one convention; the cut math stays canonical + # internally. + out_start = decanonical_start("s.seg_start", output_table) + out_end = decanonical_end("s.seg_end", output_table) + passthrough = _disjoin_passthrough(ctx, target_start, target_end, output_table) + + # Build the WITH clause one named fragment per __giql_dj_* CTE so each block + # reads on its own. The `seg_end > seg_start` guard in the final WHERE is + # belt-and-suspenders: UNION already dedupes cut positions, so LEAD cannot + # produce a zero-length segment unless it becomes UNION ALL. + ref_cte = f"__giql_dj_ref AS (SELECT * FROM {ref_from})" + tgt_cte = f"__giql_dj_tgt AS (SELECT * FROM {target_name})" + bp_cte = ( + "__giql_dj_bp AS (" + f'SELECT "{ref_chrom}" AS chrom, {bp_start} AS pos FROM __giql_dj_ref ' + "UNION " + f'SELECT "{ref_chrom}" AS chrom, {bp_end} AS pos FROM __giql_dj_ref)' + ) + # #153: alias all four columns in EVERY UNION branch. The output names already + # come from branch 1, so this is behaviour-preserving on DuckDB and makes each + # branch's projection internally unique for strict engines (DataFusion), where + # an unaliased t."end" would otherwise collide with the end-cut t."end". + cuts_cte = ( + "__giql_dj_cuts AS (" + f'SELECT t."{target_chrom}" AS kc, t."{target_start}" AS ks, ' + f't."{target_end}" AS ke, {t_start} AS pos FROM __giql_dj_tgt AS t ' + "UNION " + f'SELECT t."{target_chrom}" AS kc, t."{target_start}" AS ks, ' + f't."{target_end}" AS ke, {t_end} AS pos FROM __giql_dj_tgt AS t ' + "UNION " + f'SELECT t."{target_chrom}" AS kc, t."{target_start}" AS ks, ' + f't."{target_end}" AS ke, bp.pos AS pos ' + "FROM __giql_dj_tgt AS t JOIN __giql_dj_bp AS bp " + f"ON bp.chrom = {t_chrom} AND bp.pos > {t_start} " + f"AND bp.pos < {t_end})" + ) + segs_cte = ( + "__giql_dj_segs AS (" + "SELECT kc, ks, ke, pos AS seg_start, " + "LEAD(pos) OVER (PARTITION BY kc, ks, ke ORDER BY pos) AS seg_end " + "FROM __giql_dj_cuts)" + ) + # In self-reference mode the coverage EXISTS is provably always true: every + # emitted segment lies inside its parent target row, and that row is itself a + # member of the reference set. Skip the clause so the planner does not waste + # work on a no-op semi-join. The __giql_dj_ref CTE itself stays live because + # __giql_dj_bp still draws breakpoints from it. + where_clauses = ["s.seg_end IS NOT NULL", "s.seg_end > s.seg_start"] + if not is_self_reference: + where_clauses.append( + f'EXISTS (SELECT 1 FROM __giql_dj_ref AS r WHERE r."{ref_chrom}" = s.kc ' + f"AND {r_start} <= s.seg_start AND {r_end} > s.seg_start)" + ) + where_sql = " AND ".join(where_clauses) + final_select = ( + f"SELECT {passthrough}, s.kc AS disjoin_chrom, " + f"{out_start} AS disjoin_start, " + f"{out_end} AS disjoin_end FROM __giql_dj_tgt AS t " + f'JOIN __giql_dj_segs AS s ON t."{target_chrom}" = s.kc ' + f'AND t."{target_start}" = s.ks AND t."{target_end}" = s.ke ' + f"WHERE {where_sql}" + ) + return ( + f"(WITH {ref_cte}, {tgt_cte}, {bp_cte}, " + f"{cuts_cte}, {segs_cte} {final_select})" + ) + + +def _disjoin_passthrough( + ctx: ExpansionContext, + target_start: str, + target_end: str, + output_table: Table | None, +) -> str: + """Project the target's full row, de-canonicalizing the interval columns. + + When the target's declared encoding is canonical 0-based half-open the row + passes through as a plain ``t.*`` — the identity fast path, portable on every + engine. When it is non-canonical the interval columns (canonical inside + ``__giql_dj_tgt``) are de-canonicalized back into that encoding: + + * ``t.* REPLACE (...)`` on a target that supports it (``supports_star_replace`` + — DuckDB), substituting start/end **in place**; + * the portable ``t.* EXCEPT (start, end), , `` form otherwise + (the generic baseline / DataFusion family), which every ``* EXCEPT``-capable + engine plans. This form is **not** SQL-92 and is **not DuckDB-runnable**: it + requires ``SELECT * EXCEPT`` support. + + Column-order divergence (issue #143) + ------------------------------------ + The two forms produce row-equivalent but **not column-order-equivalent** + output. ``* REPLACE`` keeps the original column order; ``* EXCEPT`` drops + start/end and **re-appends** the recomputed columns at the end of the + projection, so they move to the right of any passthrough column on the + generic / DataFusion path. A ``SELECT *`` over a non-canonical DISJOIN is + therefore order-divergent across targets — only the row set agrees. Select + explicit columns when the cross-target column order must match. + """ + if output_table is None or ( + output_table.coordinate_system == "0based" + and output_table.interval_type == "half_open" + ): + return "t.*" + pt_start = decanonical_start(f't."{target_start}"', output_table) + pt_end = decanonical_end(f't."{target_end}"', output_table) + if ctx.capabilities.supports_star_replace: + return ( + f't.* REPLACE ({pt_start} AS "{target_start}", {pt_end} AS "{target_end}")' + ) + # Portable substitution: drop the two interval columns from the star and + # re-project them recomputed. EXCEPT removes them from the row, the trailing + # projections add them back in the target's encoding under their own names. + return ( + f't.* EXCEPT ("{target_start}", "{target_end}"), ' + f'{pt_start} AS "{target_start}", {pt_end} AS "{target_end}"' + ) + + +def _disjoin_output_encoding( + resolution: OperatorResolution | None, target_ref: ResolvedRef +) -> Table | None: + """Return the target's declared encoding for DISJOIN's output round-trip. + + ``CanonicalizeCoordinates`` (pass 2) records the original + :class:`~giql.table.Table` on the resolution when it wraps a non-canonical + target (blanking the slot's own ``table``). For an unwrapped target — a + canonical registered table, or any target when the pass did not run — the + slot's own ``table`` carries the (identity) encoding. + """ + if isinstance(resolution, OperatorResolution): + preserved = resolution.output_tables.get("this") + if preserved is not None: + return preserved + return target_ref.table + + +def _disjoin_resolution( + node: GIQLDisjoin, resolution: OperatorResolution | None +) -> tuple[ResolvedRef, ResolvedRef, str]: + """Unpack the DISJOIN resolution attached by ResolveOperatorRefs (pass 1). + + Returns ``(target_ref, ref, ref_from)`` where ``ref_from`` is the text + following ``FROM`` inside the ``__giql_dj_ref`` CTE. A subquery reference + carries no name, so it is rendered from the AST node as an aliased derived + table; registered tables and CTEs are selected from by name. + + The resolver pass deliberately leaves unresolvable slots unresolved; for + those, and for a resolved name using the reserved ``__giql_dj_`` prefix, this + re-raises DISJOIN's historical diagnostics verbatim. + """ + target_ref = ( + resolution.slot("this") + if isinstance(resolution, OperatorResolution) + else None + ) + + # An unresolved target means it is not a registered table. + if target_ref is None: + target = node.this + if isinstance(target, exp.Table): + target_name = target.name + elif isinstance(target, exp.Column): + target_name = target.table if target.table else str(target.this) + else: + target_name = str(target) + raise ValueError( + f"Target table '{target_name}' not found in tables. " + "Register the table before transpiling." + ) + + # The __giql_dj_ prefix names the operator's internal CTEs; a target table + # using it would collide with them. + if target_ref.name.startswith(DJ_PREFIX): + raise ValueError( + f"DISJOIN target {target_ref.name!r} uses the reserved " + "'__giql_dj_' prefix, which names the operator's internal " + "CTEs. Rename the table." + ) + + reference = node.args.get("reference") + ref = resolution.slot("reference") + if ref is not None: + if ref.kind == "subquery": + # Serializing the subquery with sqlglot's stock generator is safe + # because nested GIQL operators expand deepest-first: any GIQL + # construct inside this operand is already + # rewritten to standard AST by the time this outer DISJOIN expands, + # so the stock generator only ever round-trips plain SQL here. + ref_sql = reference.sql(dialect=GIQLDialect) + return target_ref, ref, f"{ref_sql} AS __giql_dj_rs" + return target_ref, ref, ref.name + + # Unresolved reference: re-classify it and raise the matching historical + # diagnostic. An omitted reference always resolves, so reference is non-None. + if not isinstance(reference, (exp.Table, exp.Column, exp.Identifier)): + raise ValueError( + "DISJOIN reference must be a table name, a CTE, or a " + f"(SELECT ...) subquery; got {type(reference).__name__}: " + f"{reference}" + ) + ref_name = reference.name + if ref_name.startswith(DJ_PREFIX): + raise ValueError( + f"DISJOIN reference {ref_name!r} uses the reserved " + "'__giql_dj_' prefix, which names the operator's internal " + "CTEs. Rename the reference relation." + ) + raise ValueError( + f"DISJOIN reference {ref_name!r} is neither a registered table " + "nor a CTE defined in this query. Register the table or define " + "the CTE before transpiling." + ) diff --git a/src/giql/expanders/distance.py b/src/giql/expanders/distance.py new file mode 100644 index 0000000..dd6f6c0 --- /dev/null +++ b/src/giql/expanders/distance.py @@ -0,0 +1,331 @@ +"""The generic DISTANCE operator expander (epic #137, step / issue #140). + +DISTANCE is the proof-of-concept that validates the expander protocol, the +registry dispatch, and the cross-target result-oracle workflow before the +harder operators migrate. It is the simplest operator — a single ``CASE`` +expression with no joins, CTEs, or per-target divergence — so a single +*generic* expander registered for :class:`~giql.targets.GenericTarget` serves +every target. DISTANCE emits identical SQL on DuckDB, DataFusion, and the +generic baseline, so no per-target override is needed. + +The CASE this expander builds matches the string builder +:func:`giql.expanders._distance.generate_distance_case` exactly +(bedtools ``closest -d`` semantics): overlapping intervals report ``0``, +book-ended (adjacent) intervals report ``1``, and a raw half-open gap of ``N`` +bases reports ``N + 1``. The ``+ 1`` is applied to the absolute gap magnitude +before any directional sign, so a downstream book-ended pair reports ``+1`` and +an upstream one ``-1`` in signed mode. There are four shapes — the cartesian +product of unsigned/signed and non-stranded/stranded — preserved verbatim from +the legacy emitter. + +Because the returned CASE is reserialized by the active target's serializer +(rather than spliced in as a raw string, as the legacy ``giqldistance_sql`` +emitter did), the emitted text changes cosmetically — most visibly ``!=`` +renders as the SQL-standard ``<>``. The two are semantically identical. +""" + +from __future__ import annotations + +from sqlglot import exp +from sqlglot import parse_one + +from giql.dialect import GIQLDialect +from giql.expander import ExpansionContext +from giql.expander import register +from giql.expanders._params import coerce_bool_param +from giql.expressions import GIQLDistance +from giql.resolver import ResolvedColumn +from giql.targets import GenericTarget + +__all__ = ["expand_distance"] + + +def _frag(fragment: str) -> exp.Expression: + """Parse one canonicalized SQL fragment into AST. + + The pass-1 :class:`~giql.resolver.ResolvedColumn` endpoints are SQL string + fragments (e.g. ``a."end"``, or ``'chr1'`` for a literal range) already + canonicalized in place by pass 2, so they are parsed — not rebuilt — back + into AST under the GIQL dialect. + + Parameters + ---------- + fragment : str + A canonicalized SQL fragment (a column reference or literal). + + Returns + ------- + exp.Expression + The parsed fragment as a sqlglot AST node. + """ + return parse_one(fragment, dialect=GIQLDialect) + + +def _gap(minuend: str, subtrahend: str) -> exp.Expression: + """Build ``(minuend - subtrahend + 1)`` — the bedtools-parity gap magnitude. + + Mirrors the legacy ``({start} - {end} + 1)`` fragment: the ``+ 1`` lifts a + book-ended (adjacent) pair from a raw half-open gap of ``0`` to a reported + distance of ``1``. + + Parameters + ---------- + minuend : str + The SQL fragment subtracted *from* (the left operand of the ``-``). + subtrahend : str + The SQL fragment subtracted (the right operand of the ``-``). + + Returns + ------- + exp.Expression + The parenthesized ``(minuend - subtrahend + 1)`` AST. + """ + diff = exp.Sub(this=_frag(minuend), expression=_frag(subtrahend)) + return exp.paren(exp.Add(this=diff, expression=exp.Literal.number(1))) + + +def _operand(ctx: ExpansionContext, arg: str, position: str) -> ResolvedColumn: + """Return the resolved column for one DISTANCE interval operand. + + Reads the pass-1 metadata attached to the node. A deferred operand (a + literal range, or an unqualified column the resolver could not resolve) has + no column attached; this raises the historical literal-range diagnostic so + the public error contract is preserved. + + Parameters + ---------- + ctx : ExpansionContext + The expansion context carrying the node's pass-1 resolution. + arg : str + The operand slot key (``"this"`` or ``"expression"``). + position : str + Human-readable operand position (``"first"`` / ``"second"``) for the + diagnostic message. + + Returns + ------- + ResolvedColumn + The resolved column metadata for the operand. + + Raises + ------ + ValueError + If the operand was deferred (a literal range or unresolved column). + """ + # TODO: this read-required-column-or-raise pattern is duplicated across + # expanders; hoist it to a shared ``ExpansionContext.require_column`` helper + # once a second expander needs it. + resolution = ctx.resolution + if resolution is not None: + resolved = resolution.column(arg) + if resolved is not None: + return resolved + raise ValueError(f"Literal range as {position} argument not yet supported") + + +def _unsigned_distance(col_a: ResolvedColumn, col_b: ResolvedColumn) -> exp.Expression: + """Branch 1: unsigned (absolute) non-stranded distance, returning ``|gap| + 1``.""" + return _wrap_overlap_case( + col_a, + col_b, + downstream=_gap(col_b.start, col_a.end), + upstream=_gap(col_a.start, col_b.end), + ) + + +def _signed_distance(col_a: ResolvedColumn, col_b: ResolvedColumn) -> exp.Expression: + """Branch 2: signed non-stranded distance. + + ``+`` downstream (B after A), ``-`` upstream (B before A). + """ + return _wrap_overlap_case( + col_a, + col_b, + downstream=_gap(col_b.start, col_a.end), + upstream=exp.Neg(this=_gap(col_a.start, col_b.end)), + ) + + +def _stranded_distance( + col_a: ResolvedColumn, col_b: ResolvedColumn, signed: bool +) -> exp.Expression: + """Branches 3 & 4: stranded distance, flipping sign on A's ``-`` strand. + + The downstream and upstream gaps each become a nested ``CASE`` keyed on + ``strand_a``. The ``signed`` flag additionally layers the directional sign + on top of the strand flip, exactly as the legacy emitter's two stranded + branches do. + """ + strand_a = col_a.strand + strand_b = col_b.strand + assert strand_a is not None and strand_b is not None # gated by caller + + down_gap = _gap(col_b.start, col_a.end) + up_gap = _gap(col_a.start, col_b.end) + + # Downstream (B after A) is identical across the signed and unsigned arms: + # positive by default, flipped negative on A's '-' strand. Only *upstream* + # differs between the arms, so hoist downstream and compute upstream per arm. + downstream = _strand_flip_case( + strand_a, neg=exp.Neg(this=down_gap), pos=down_gap.copy() + ) + if signed: + # Stranded + signed: upstream (B before A) is negative by default but + # flips positive on A's '-' strand (the directional sign layered on top + # of the strand flip). + upstream = _strand_flip_case( + strand_a, neg=up_gap, pos=exp.Neg(this=up_gap.copy()) + ) + else: + # Stranded but not signed: upstream carries the strand flip only. + upstream = _strand_flip_case( + strand_a, neg=exp.Neg(this=up_gap), pos=up_gap.copy() + ) + + case = _wrap_overlap_case(col_a, col_b, downstream=downstream, upstream=upstream) + # Prepend the strand-validity guards ahead of the overlap guards. The WHEN + # order matters: chrom mismatch, then strand NULL/'.'/'?', then overlap. + return _prepend_strand_guards(case, strand_a, strand_b) + + +def _strand_flip_case( + strand_a: str, neg: exp.Expression, pos: exp.Expression +) -> exp.Expression: + """Build ``CASE WHEN strand_a = '-' THEN neg ELSE pos END``.""" + return ( + exp.Case() + .when(exp.EQ(this=_frag(strand_a), expression=exp.Literal.string("-")), neg) + .else_(pos) + ) + + +def _wrap_overlap_case( + col_a: ResolvedColumn, + col_b: ResolvedColumn, + downstream: exp.Expression, + upstream: exp.Expression, +) -> exp.Expression: + """Build the shared distance CASE skeleton common to all four branches. + + ``CASE WHEN chrom_a != chrom_b THEN NULL WHEN THEN 0 WHEN + end_a <= start_b THEN ELSE END``. + + Parameters + ---------- + col_a, col_b : ResolvedColumn + The resolved A and B interval operands. + downstream : exp.Expression + The B-after-A gap expression (already carrying any sign / strand flip). + upstream : exp.Expression + The B-before-A gap expression (already carrying any sign / strand flip). + + Returns + ------- + exp.Expression + The assembled distance ``CASE`` expression. + """ + overlap = exp.and_( + exp.LT(this=_frag(col_a.start), expression=_frag(col_b.end)), + exp.GT(this=_frag(col_a.end), expression=_frag(col_b.start)), + ) + chrom_mismatch = exp.NEQ(this=_frag(col_a.chrom), expression=_frag(col_b.chrom)) + end_a_le_start_b = exp.LTE(this=_frag(col_a.end), expression=_frag(col_b.start)) + return ( + exp.Case() + .when(chrom_mismatch, exp.Null()) + .when(overlap, exp.Literal.number(0)) + .when(end_a_le_start_b, downstream) + .else_(upstream) + ) + + +def _prepend_strand_guards( + case: exp.Case, strand_a: str, strand_b: str +) -> exp.Case: + """Insert the strand-validity WHEN guards after the chrom guard. + + Distance is undefined for an unstranded (``'.'``/``'?'``) or missing strand, + matching the legacy emitter. + + Parameters + ---------- + case : exp.Case + The overlap/gap ``CASE`` to prepend the strand guards onto (mutated in + place). + strand_a, strand_b : str + The A and B strand-column SQL fragments. + + Returns + ------- + exp.Case + The same *case*, whose WHEN order is now: chrom mismatch -> NULL, either + strand NULL -> NULL, ``strand_a`` is ``'.'``/``'?'`` -> NULL, + ``strand_b`` is ``'.'``/``'?'`` -> NULL, then the original overlap/gap + branches. + """ + sa = _frag(strand_a) + sb = _frag(strand_b) + null_guard = exp.condition(exp.Is(this=sa, expression=exp.Null())).or_( + exp.Is(this=sb.copy(), expression=exp.Null()) + ) + a_unstranded = exp.condition( + exp.EQ(this=sa.copy(), expression=exp.Literal.string(".")) + ).or_(exp.EQ(this=sa.copy(), expression=exp.Literal.string("?"))) + b_unstranded = exp.condition( + exp.EQ(this=sb.copy(), expression=exp.Literal.string(".")) + ).or_(exp.EQ(this=sb.copy(), expression=exp.Literal.string("?"))) + + guards = [ + exp.If(this=null_guard, true=exp.Null()), + exp.If(this=a_unstranded, true=exp.Null()), + exp.If(this=b_unstranded, true=exp.Null()), + ] + # The existing WHENs keep their order; the strand guards slot in right after + # the leading chrom-mismatch guard. + existing = case.args["ifs"] + case.set("ifs", existing[:1] + guards + existing[1:]) + return case + + +# KEEP IN SYNC: this expander (AST) and the string builder +# ``giql.expanders._distance.generate_distance_case`` build the *same* distance +# CASE by two routes. The string form is retained because NEAREST assembles its +# ORDER BY / filter math string-first and splices the CASE in. Any change to the +# distance math here must be mirrored there (and vice versa). The parity test in +# tests/test_distance_udf.py guards the drift. +@register(GenericTarget, GIQLDistance) +def expand_distance(node: exp.Expression, ctx: ExpansionContext) -> exp.Expression: + """Expand a ``GIQLDistance`` node into a standard SQL ``CASE`` expression. + + The portable expander registered for every target. Reads ``stranded`` / + ``signed`` from the node and the pass-1 resolved interval operands, then + builds the matching one of the four CASE shapes + (unsigned/signed x non-stranded/stranded). + + Parameters + ---------- + node : exp.Expression + The ``GIQLDistance`` operator node being expanded. + ctx : ExpansionContext + The expansion context carrying the node's pass-1 resolution. + + Returns + ------- + exp.Expression + The distance ``CASE`` that replaces the operator node and is rendered by + the active target's serializer. + """ + stranded = coerce_bool_param(node.args.get("stranded")) + signed = coerce_bool_param(node.args.get("signed")) + + col_a = _operand(ctx, "this", "first") + col_b = _operand(ctx, "expression", "second") + + # Strand columns are consumed only in stranded mode, and only when both + # operands actually carry a strand fragment — mirroring the legacy emitter's + # `strand_a is None or strand_b is None` fall-through to the unstranded path. + if stranded and col_a.strand is not None and col_b.strand is not None: + return _stranded_distance(col_a, col_b, signed=signed) + if signed: + return _signed_distance(col_a, col_b) + return _unsigned_distance(col_a, col_b) diff --git a/src/giql/expanders/intersects.py b/src/giql/expanders/intersects.py new file mode 100644 index 0000000..938587b --- /dev/null +++ b/src/giql/expanders/intersects.py @@ -0,0 +1,270 @@ +"""Generic expanders for the spatial predicates and set predicates (epic #137). + +Migrates INTERSECTS / CONTAINS / WITHIN and the ``ANY`` / ``ALL`` set predicates +off the legacy ``*_sql`` generator emitters and onto the operator-expander +registry. Each expander turns one predicate node +into standard sqlglot AST built from the pass-1 :class:`~giql.resolver.ResolvedColumn` +metadata (already canonicalized to 0-based half-open by pass 2), so the emitted SQL +is byte-identical to the strings the legacy emitter produced. + +These are *node-local* predicate rewrites: an INTERSECTS / CONTAINS / WITHIN node +expands to a boolean ``(chrom = ... AND start < ... AND end > ...)`` expression that +replaces it in place. For a column-to-column INTERSECTS *join* this in-place +expansion IS the join plan on every target but DuckDB: the boolean overlap +predicate is a plain ``ON`` condition the engine plans as a range join (a hash join +keyed on ``chrom`` with the position inequalities as a residual filter), correct for +both inner and outer joins. (The generic binned equi-join was dropped in favor of +this naive predicate — #167.) + +DuckDB overrides the column-to-column *join* case with its per-chromosome IEJoin +plan, registered as the ``(DuckDBTarget, Intersects)`` expander in +:mod:`giql.expanders.intersects_duckdb` (#169) — a pass-3 override, not a pre-pass. +That override reuses this module's :func:`_expand_spatial_op` for every shape it +declines (LEFT/RIGHT/FULL, self-joins, multiple INTERSECTS, …) and for literal-range +or residual INTERSECTS *predicates*, so the naive predicate here is the universal +fallback on DuckDB too. + +The INTERSECTS / CONTAINS / WITHIN / set-predicate expanders are registered only for +:class:`~giql.targets.GenericTarget`: spatial-predicate *emission* is portable +SQL-92 and does not vary by engine, so one generic expander covers every target via +the registry's ``(generic, op)`` fallback (the DuckDB IEJoin override is the sole +target-specific INTERSECTS entry, and it delegates back here). +""" + +from __future__ import annotations + +from sqlglot import exp +from sqlglot import maybe_parse + +from giql.dialect import GIQLDialect +from giql.expander import ExpansionContext +from giql.expander import register +from giql.expressions import Contains +from giql.expressions import Intersects +from giql.expressions import SpatialSetPredicate +from giql.expressions import Within +from giql.range_parser import ParsedRange +from giql.range_parser import RangeParser +from giql.resolver import ResolvedColumn +from giql.targets import GenericTarget + + +def _fragment(fragment: str) -> exp.Expression: + """Parse a resolved SQL fragment (e.g. ``a."end"`` / ``'chr1'``) into AST. + + The pass-1 :class:`~giql.resolver.ResolvedColumn` carries column references as + pre-canonicalized SQL string fragments; parse them through the GIQL dialect so + the rebuilt predicate reserializes identically to the legacy emitter's string. + """ + parsed = maybe_parse(fragment, dialect=GIQLDialect) + if parsed is None: + # maybe_parse returns None only for an empty/None input; a ResolvedColumn + # fragment is never empty, so this is an internal invariant violation. + raise ValueError(f"Could not parse resolved column fragment: {fragment!r}") + return parsed + + +def _predicate_column(ctx: ExpansionContext, arg: str) -> ResolvedColumn: + """Return the :class:`ResolvedColumn` for predicate operand *arg*. + + Mirrors the legacy ``_predicate_operand`` emitter helper: the + expander consumes only the pass-1 resolution; a missing column means pass 1 did + not run (an internal invariant violation), so raise the historical message. + """ + resolution = ctx.resolution + if resolution is not None: + resolved = resolution.column(arg) + if resolved is not None: + return resolved + raise ValueError( + f"Spatial predicate operand {arg!r} was not resolved; run the " + "ResolveOperatorRefs pass (transpile pipeline) before generation." + ) + + +def _range_predicate( + column: ResolvedColumn, parsed: ParsedRange, op_type: str +) -> exp.Expression: + """Build the boolean AST for ``column ``. + + Reproduces the legacy ``_generate_range_predicate`` emitter as AST. The + column fragments are already canonical 0-based half-open (pass 2); the parsed + range is canonicalized by the caller. Returns a parenthesized boolean. + """ + chrom = _fragment(column.chrom) + start = _fragment(column.start) + end = _fragment(column.end) + chrom_lit = exp.Literal.string(parsed.chromosome) + r_start = exp.Literal.number(parsed.start) + r_end = exp.Literal.number(parsed.end) + + if op_type == "intersects": + # Ranges overlap if: start1 < end2 AND end1 > start2 + cond = exp.and_( + exp.EQ(this=chrom, expression=chrom_lit), + exp.LT(this=start, expression=r_end), + exp.GT(this=end, expression=r_start), + ) + elif op_type == "contains": + if parsed.end == parsed.start + 1: + # Point query: start1 <= point < end1 + cond = exp.and_( + exp.EQ(this=chrom, expression=chrom_lit), + exp.LTE(this=start, expression=r_start), + exp.GT(this=end, expression=r_start), + ) + else: + # Range query: start1 <= start2 AND end1 >= end2 + cond = exp.and_( + exp.EQ(this=chrom, expression=chrom_lit), + exp.LTE(this=start, expression=r_start), + exp.GTE(this=end, expression=r_end), + ) + elif op_type == "within": + # left within right: start1 >= start2 AND end1 <= end2 + cond = exp.and_( + exp.EQ(this=chrom, expression=chrom_lit), + exp.GTE(this=start, expression=r_start), + exp.LTE(this=end, expression=r_end), + ) + else: + raise ValueError(f"Unknown spatial op_type: {op_type!r}") + + return exp.paren(cond) + + +def _column_join( + left: ResolvedColumn, right: ResolvedColumn, op_type: str +) -> exp.Expression: + """Build the boolean AST for a column-to-column spatial predicate. + + Reproduces the legacy ``_generate_column_join`` emitter as AST. Both + operands' fragments are pre-canonicalized (pass 2). Returns a parenthesized + boolean. + """ + l_chrom, r_chrom = _fragment(left.chrom), _fragment(right.chrom) + l_start, r_start = _fragment(left.start), _fragment(right.start) + l_end, r_end = _fragment(left.end), _fragment(right.end) + + if op_type == "intersects": + cond = exp.and_( + exp.EQ(this=l_chrom, expression=r_chrom), + exp.LT(this=l_start, expression=r_end), + exp.GT(this=l_end, expression=r_start), + ) + elif op_type == "contains": + cond = exp.and_( + exp.EQ(this=l_chrom, expression=r_chrom), + exp.LTE(this=l_start, expression=r_start), + exp.GTE(this=l_end, expression=r_end), + ) + elif op_type == "within": + cond = exp.and_( + exp.EQ(this=l_chrom, expression=r_chrom), + exp.GTE(this=l_start, expression=r_start), + exp.LTE(this=l_end, expression=r_end), + ) + else: + raise ValueError(f"Unknown spatial op_type: {op_type!r}") + + return exp.paren(cond) + + +#: Meta key marking an AST node as the expansion of a GIQL spatial predicate +#: (INTERSECTS / CONTAINS / WITHIN / ``ANY`` / ``ALL``). The DuckDB IEJoin override +#: (:mod:`giql.expanders.intersects_duckdb`) reads it to decline its per-chromosome +#: rewrite when a *sibling* spatial predicate shares the query. Because pass 3 may +#: expand that sibling to plain comparison SQL before the join ``INTERSECTS`` node +#: is visited, the override cannot rely on finding an un-expanded GIQL node — the +#: tag survives expansion so the sibling stays detectable, preserving the former +#: pre-pass's "decline when a residual spatial predicate is present" invariant (#169). +SPATIAL_PREDICATE_META = "__giql_spatial_predicate" + + +def _tag_spatial(expr: exp.Expression) -> exp.Expression: + """Mark *expr* as a spatial-predicate expansion (:data:`SPATIAL_PREDICATE_META`).""" + expr.meta[SPATIAL_PREDICATE_META] = True + return expr + + +def _expand_spatial_op( + node: exp.Expression, ctx: ExpansionContext, op_type: str +) -> exp.Expression: + """Expand one INTERSECTS / CONTAINS / WITHIN node to a boolean predicate. + + Dispatches on the right operand exactly as the legacy emitter did: the + presence of a resolved right *column* — keyed off + ``ctx.resolution.column("expression")``, the slot pass 1 attaches a + :class:`ResolvedColumn` to when the right operand is a column reference — + selects the column-to-column path; its absence means the right operand is a + literal range, parsed in place. The result carries the + :data:`SPATIAL_PREDICATE_META` tag so a sibling DuckDB IEJoin join override can + still recognize this expansion after the fact (#169). + """ + resolution = ctx.resolution + right_column = resolution.column("expression") if resolution is not None else None + left = _predicate_column(ctx, "this") + + if right_column is not None: + return _tag_spatial(_column_join(left, right_column, op_type)) + + # Literal range string (e.g. interval INTERSECTS 'chr1:1000-2000'). Reproduce + # the legacy emitter's parse-and-wrap-error behavior verbatim: any parse + # failure (including the RangeParser's own ValueError) is wrapped in the + # historical "Could not parse genomic range" message. + right_expr = node.args.get("expression") + raw = right_expr.sql(dialect=GIQLDialect) if right_expr is not None else "" + try: + range_str = raw.strip("'\"") + parsed = RangeParser.parse(range_str).to_zero_based_half_open() + return _tag_spatial(_range_predicate(left, parsed, op_type)) + except Exception as e: + raise ValueError(f"Could not parse genomic range: {raw}. Error: {e}") from e + + +@register(GenericTarget, Intersects) +def expand_intersects(node: exp.Expression, ctx: ExpansionContext) -> exp.Expression: + """Expand an INTERSECTS predicate to standard boolean SQL AST.""" + return _expand_spatial_op(node, ctx, "intersects") + + +@register(GenericTarget, Contains) +def expand_contains(node: exp.Expression, ctx: ExpansionContext) -> exp.Expression: + """Expand a CONTAINS predicate to standard boolean SQL AST.""" + return _expand_spatial_op(node, ctx, "contains") + + +@register(GenericTarget, Within) +def expand_within(node: exp.Expression, ctx: ExpansionContext) -> exp.Expression: + """Expand a WITHIN predicate to standard boolean SQL AST.""" + return _expand_spatial_op(node, ctx, "within") + + +@register(GenericTarget, SpatialSetPredicate) +def expand_spatial_set(node: exp.Expression, ctx: ExpansionContext) -> exp.Expression: + """Expand a quantified set predicate (``ANY`` / ``ALL``) to boolean SQL AST. + + Reproduces the legacy ``_generate_spatial_set`` emitter: the single left + column is compared against every literal range, and the per-range conditions + are OR-combined for ``ANY`` / AND-combined for ``ALL``, all wrapped in one + outer paren. + """ + operator = node.args["operator"] + quantifier = node.args["quantifier"] + ranges = node.args["ranges"] + + column = _predicate_column(ctx, "this") + op_type = operator.lower() + + conditions: list[exp.Expression] = [] + for range_expr in ranges: + range_str = range_expr.sql(dialect=GIQLDialect).strip("'\"") + parsed = RangeParser.parse(range_str).to_zero_based_half_open() + conditions.append(_range_predicate(column, parsed, op_type)) + + if quantifier.upper() == "ANY": + combined = exp.or_(*conditions) + else: + combined = exp.and_(*conditions) + + return _tag_spatial(exp.paren(combined)) diff --git a/src/giql/expanders/intersects_duckdb.py b/src/giql/expanders/intersects_duckdb.py new file mode 100644 index 0000000..f5bbf09 --- /dev/null +++ b/src/giql/expanders/intersects_duckdb.py @@ -0,0 +1,1715 @@ +"""The DuckDB IEJoin column-to-column INTERSECTS join override (epic #137, #169). + +The registered ``(DuckDBTarget, Intersects)`` override expander for DuckDB. +:class:`IntersectsDuckDBIEJoinTransformer` rewrites a column-to-column INTERSECTS +*join* into a per-chromosome dynamic-SQL pattern DuckDB plans through its +range-join family (``IE_JOIN`` / ``PIECEWISE_MERGE_JOIN``); +:func:`expand_intersects_duckdb` wires it into the operator-expander registry so +pass 3 (:class:`giql.expander.ExpandOperators`) dispatches DuckDB's ``Intersects`` +nodes here. For the shapes it declines, and for every literal-range or residual +``Intersects`` predicate, the override defers to the shared naive overlap +predicate (:func:`giql.expanders.intersects._expand_spatial_op`) that +``dialect=None`` / ``"datafusion"`` emit on every target — a plain ``ON`` condition +the engine plans as a range join (#167). + +Because the IEJoin rewrite emits a *whole-query* multi-statement string +(``SET VARIABLE ...; SELECT ... FROM query(getvariable(...))``) rather than a +node-local expression, :func:`expand_intersects_duckdb` produces it through the +query-level :meth:`~giql.expander.ExpansionContext.add_statement_finalizer` seam, +wrapping the built string in an :class:`sqlglot.expressions.Command` so it +serializes verbatim through the stock per-target serializer. This relocates the +former capability-gated pre-pass IEJoin transformer into the registry, mirroring +the CLUSTER / MERGE relocation (#144); INTERSECTS join strategy is now fully +registry-driven and target-overridable. The generic binned equi-join transformer +was dropped in favor of the naive predicate earlier (#167). +""" + +from dataclasses import dataclass +from typing import ClassVar +from typing import Literal + +from sqlglot import exp + +from giql.canonical import canonical_end +from giql.canonical import canonical_start +from giql.constants import DEFAULT_CHROM_COL +from giql.constants import DEFAULT_END_COL +from giql.constants import DEFAULT_START_COL +from giql.expander import ExpansionContext +from giql.expander import register +from giql.expanders._per_chrom import CHROM_LITERAL +from giql.expanders._per_chrom import dynamic_relation +from giql.expanders._per_chrom import new_variable_name +from giql.expanders._per_chrom import set_variable_statement +from giql.expanders.intersects import SPATIAL_PREDICATE_META +from giql.expanders.intersects import _expand_spatial_op +from giql.expressions import Contains +from giql.expressions import Intersects +from giql.expressions import SpatialSetPredicate +from giql.expressions import Within +from giql.table import Tables +from giql.targets import DuckDBTarget + + +class _UnqualifiedProjectionError(Exception): + """Signal an unrecoverable projection-resolution failure in the DuckDB IEJoin path. + + Raised by :class:`IntersectsDuckDBIEJoinTransformer` when a query engages + the dialect path but contains a SELECT-list shape the rebuild cannot + attribute to a join side: an unqualified column, an unknown table + qualifier, an aggregate over an unqualified column, or a right-side + reference under a SEMI / ANTI join, plus an extras predicate that + references an unqualified or unknown-aliased column. The unknown-table, + aggregate-of-unqualified, and right-side cases are also rejected by the + naive plan; a *bare* unqualified column may in fact be naive-valid when it + is unambiguous, but the dialect has no live schema to attribute it to a + side, so it surfaces a clean transpile-time error rather than guessing. + Caught and re-raised as :class:`ValueError` at the dialect's public + boundary. + + Shapes the naive plan *accepts* but the IEJoin cannot rebuild do NOT raise + — star projections (#202), and expressions / window aggregates / ``FILTER`` + clauses / scalar subqueries / aggregate-nested stars (#204, #205) — they + signal :class:`_DeclineIEJoin` and fall back to the naive plan instead. + """ + + +class _DeclineIEJoin(Exception): + """Signal that the IEJoin should decline a valid query to the naive plan. + + Distinct from :class:`_UnqualifiedProjectionError`: this is not a user + error but a shape the naive-predicate plan compiles for free while the + IEJoin projection rebuild cannot express it (an expression, a window + aggregate, a ``FILTER`` clause, a scalar subquery, an aggregate nested in + an expression, or a star nested in an aggregate argument — #204, #205). + Caught in :meth:`IntersectsDuckDBIEJoinTransformer.transform_to_sql`, which + returns ``None`` so the generic naive-predicate plan handles the query, + keeping ``dialect="duckdb"`` consistent with every other backend. + """ + + +# These predicates support :class:`IntersectsDuckDBIEJoinTransformer`; they live +# at module scope rather than on the class so a future target-specific INTERSECTS +# override can reuse them without importing the transformer. Style §2's strict +# "functions after classes" ordering is relaxed here under the topic-interleaved +# exception. + + +def _is_column_intersects(node: Intersects) -> bool: + """Return True if *node* is column-to-column :class:`Intersects`. + + Both operands must be table-qualified column references. + """ + return ( + isinstance(node.this, exp.Column) + and bool(node.this.table) + and isinstance(node.expression, exp.Column) + and bool(node.expression.table) + ) + + +def _find_column_intersects_in(expr: exp.Expression) -> Intersects | None: + """Return the first column-to-column :class:`Intersects` in *expr*, or ``None``.""" + return next( + (n for n in expr.find_all(Intersects) if _is_column_intersects(n)), + None, + ) + + +def _normalize_alias(name: str, *, quoted: bool = False) -> str: + """Return *name* case-folded unless it was a quoted identifier. + + DuckDB (and standard SQL) treat unquoted identifiers case-insensitively. + Comparing aliases via raw Python equality rejects valid mixed-case queries + that the naive-predicate plan accepts; normalize via :py:meth:`str.casefold` so the + dialect matches DuckDB's identifier semantics. Quoted identifiers stay + case-sensitive per the SQL standard. + """ + return name if quoted else name.casefold() + + +def _has_outer_join_intersects(query: exp.Select) -> bool: + """Return True if any outer JOIN has a column-to-column :class:`Intersects`. + + Handles two surfaces the user can write the INTERSECTS under an outer + join: directly in the join's ``ON`` clause, or in the top-level + ``WHERE`` while the join keeps its ``LEFT``/``RIGHT``/``FULL`` side + modifier (``LEFT JOIN ... ON TRUE WHERE a.interval INTERSECTS + b.interval``). The dialect's inner-join IEJoin rewrite would silently + drop the outer-join's unmatched-row guarantee in the second shape, + so we fall back whenever an outer-join is present and an + INTERSECTS sits anywhere it could influence the join result. + """ + if not isinstance(query, exp.Select): + return False + outer_joins = [ + join for join in query.args.get("joins") or [] if join.args.get("side") + ] + if not outer_joins: + return False + for join in outer_joins: + on = join.args.get("on") + if on is not None and _find_column_intersects_in(on): + return True + where = query.args.get("where") + if where and _find_column_intersects_in(where.this): + return True + return False + + +def _has_left_only_join_where_intersects(query: exp.Select) -> bool: + """Return True if a SEMI/ANTI join carries its INTERSECTS in the WHERE. + + A left-only (``SEMI`` / ``ANTI``) join's right table is out of scope in + the top-level ``WHERE`` — the reference plans (``dialect=None`` / + ``"datafusion"``) reject a right-referencing ``WHERE`` predicate with a + binder error. The dialect's ``WHERE``-branch (which exists to support the + legitimate implicit comma/cross-join idiom, an ``INNER`` join whose right + table *is* in scope in the ``WHERE``) would instead relocate that + ``INTERSECTS`` into the join ``ON`` and invent anti/semi-overlap results, + diverging from every other backend (#201). Decline whenever an explicit + ``SEMI`` / ``ANTI`` join is present and a column-to-column ``INTERSECTS`` + sits in the top-level ``WHERE`` (necessarily out of scope), so the naive + plan surfaces the same binder error the reference plans raise. The + ``INTERSECTS`` search recurses (via ``find_all``), so a column-to-column + ``INTERSECTS`` nested inside a ``WHERE`` subquery also declines here; + that is a harmless over-decline (it only ever routes to the correct + naive plan) and such shapes already decline via other gates. + """ + if not isinstance(query, exp.Select): + return False + has_left_only_join = any( + (join.args.get("kind") or "").upper() in ("SEMI", "ANTI") + for join in query.args.get("joins") or [] + ) + if not has_left_only_join: + return False + where = query.args.get("where") + return bool(where and _find_column_intersects_in(where.this)) + + +def _has_star_projection(query: exp.Select) -> bool: + """Return True if the SELECT list contains any star projection. + + The IEJoin rewrite restructures the query into per-chromosome subqueries + that project **explicitly-named** ``__giql_p`` columns up through a + wrapper relation, so building that dynamic SQL requires enumerating the + star's columns at transpile time. But GIQL only knows the genomic columns + declared in the :class:`~giql.table.Table` config (chrom / start / end / + strand) — it has no view of the live table schema, so it cannot see + arbitrary base-table columns (``name`` / ``score`` / ...). Expanding + ``a.*`` here would silently drop them, narrowing the result relative to + every other backend (#202). Decline any star projection (bare ``*`` or + qualified ``t.*``, with or without a user alias) so the naive-predicate + plan expands it against DuckDB's live schema at bind time, keeping the + projection identical across backends. + """ + if not isinstance(query, exp.Select): + return False + for sel in query.expressions: + target = sel.this if isinstance(sel, exp.Alias) else sel + while isinstance(target, exp.Paren): + target = target.this + if isinstance(target, exp.Star): + return True + if isinstance(target, exp.Column) and isinstance(target.this, exp.Star): + return True + return False + + +def _projection_declines_to_naive( + target: exp.Expression, + l_alias: str, + r_alias: str, + is_left_only: bool, +) -> bool: + """Return True if *target* is a projection the IEJoin should decline to naive. + + The IEJoin projection rebuild only handles a qualified column (``a.col``) + or a plain aggregate over qualified columns (``SUM(a.score)`` / ``COUNT(*)``); + it emits explicitly-named ``__giql_p`` inner columns and reconstructs the + outer SELECT from them. Any other shape the naive-predicate plan compiles + for free — an expression (``a.start + 1``), a window aggregate, a ``FILTER`` + clause, a scalar subquery, an aggregate nested in an expression + (``COUNT(*) * 2``), or a star nested in an aggregate argument + (``COUNT(a.*)`` / ``MIN(COLUMNS(*))``) — should decline so ``dialect="duckdb"`` + stays consistent with every other backend instead of hard-erroring or, for + the aggregate-nested star, miscompiling (#204, #205). + + A projection that references a column **in its own scope** the rebuild + cannot attribute to a join side (unqualified, an unknown table, or the + right side under a SEMI / ANTI join) does NOT decline — the dispatch raises + a clean :class:`_UnqualifiedProjectionError` instead. The unknown-table and + right-side cases are rejected by the naive plan too; a bare unqualified + column may be naive-valid when unambiguous, but the dialect has no live + schema to pick a side, so a transpile-time error is more actionable than + guessing. Columns inside a nested subquery resolve against that subquery, + not the join, so they are ignored by this test. + """ + # A qualified, non-star column is rebuilt directly — never declines here. + if ( + isinstance(target, exp.Column) + and target.table + and not isinstance(target.this, exp.Star) + ): + return False + # An aggregate is decided here in full (it never falls through to the + # column scan below — an unqualified aggregate argument like ``SUM(score)`` + # is diagnosed in ``render_aggregate``, not treated as a decline). A star + # nested in an aggregate argument (``COUNT(a.*)`` / ``MIN(COLUMNS(*))`` / + # ``COUNT(a.* EXCLUDE (name))``) cannot be enumerated, so it declines + # regardless of any EXCLUDE / REPLACE modifier columns (#204); a plain + # aggregate over qualified columns (``COUNT(*)`` / ``SUM(a.score)`` / + # ``COUNT(DISTINCT a.col)``) is rebuilt directly. + if isinstance(target, exp.AggFunc): + has_star_arg = any( + isinstance(col.this, exp.Star) for col in target.find_all(exp.Column) + ) + return has_star_arg or target.find(exp.Columns) is not None + # Unsupported by the rebuild. Decline to naive unless an own-scope column + # is out of scope (a user error the naive plan rejects too — let the + # dispatch raise instead). + for col in target.find_all(exp.Column): + if isinstance(col.this, exp.Star): + continue # qualified star — declinable, not a plain column ref + # Skip a column inside a nested subquery within *target* (it resolves + # against that subquery, not the join). *target* is an ancestor of + # *col*, so walking up from ``col.parent`` reaches it; a Select / + # Subquery strictly between them is a nested scope. When ``col`` IS + # ``target`` (a bare column projection), there is no interior to walk. + in_subquery = False + node = col.parent if col is not target else target + while node is not None and node is not target: + if isinstance(node, (exp.Select, exp.Subquery)): + in_subquery = True + break + node = node.parent + if in_subquery: + continue + tbl = _normalize_alias(col.table) if col.table else "" + if tbl not in (l_alias, r_alias) or (is_left_only and tbl == r_alias): + return False + return True + + +def _strip_intersects( + expr: exp.Expression | None, intersects: Intersects +) -> exp.Expression | None: + """Return the parts of an AND tree that aren't the given :class:`Intersects`.""" + if expr is None or expr is intersects: + return None + if isinstance(expr, exp.And): + if expr.this is intersects: + return expr.expression + if expr.expression is intersects: + return expr.this + left = _strip_intersects(expr.this, intersects) + right = _strip_intersects(expr.expression, intersects) + if left is None: + return right + if right is None: + return left + return exp.And(this=left, expression=right) + return expr + + +def _count_column_intersects(query: exp.Select) -> int: + """Count column-to-column :class:`Intersects` predicates anywhere in *query*. + + Walks the full AST so the dialect's "exactly one INTERSECTS" gate cannot + be circumvented by an INTERSECTS hidden inside a subquery (e.g. inside + a SELECT-list scalar subquery or inside ``EXISTS (SELECT ...)``). The + SELECT-list / modifier-clause subquery gates in ``transform_to_sql`` + already reject those shapes, but this widened count is defense-in-depth + against future gate relaxations. + """ + if not isinstance(query, exp.Select): + return 0 + return sum(1 for n in query.find_all(Intersects) if _is_column_intersects(n)) + + +@dataclass(frozen=True) +class _IEJoinSides: + """Resolved two-sided join shape for the DuckDB IEJoin dialect path. + + Bundles the left/right user aliases (normalized via :func:`_normalize_alias`) + with the un-aliased :class:`~sqlglot.expressions.Table` AST nodes and a + ``kind`` discriminator identifying the join variant. The hardcoded inner + aliases ``a`` and ``b`` used inside each per-chromosome subquery are + exposed as :data:`left_inner_alias` and :data:`right_inner_alias`. + + The :attr:`is_left_only_join` property is True for SEMI and ANTI shapes, + which restrict the outer SELECT to left-side projections only. + """ + + left_user_alias: str + right_user_alias: str + left_table: exp.Table + right_table: exp.Table + kind: Literal["INNER", "SEMI", "ANTI"] + + left_inner_alias: ClassVar[str] = "a" + right_inner_alias: ClassVar[str] = "b" + + @property + def is_left_only_join(self) -> bool: + """True when the join's output only references the left side (SEMI / ANTI).""" + return self.kind in ("SEMI", "ANTI") + + @classmethod + def from_intersects( + cls, + from_table: exp.Table, + join: exp.Join, + intersects: Intersects, + ) -> "_IEJoinSides | None": + """Resolve sides, run the FROM-side orientation swap, normalize aliases. + + Returns ``None`` when the join's table operand isn't a base + :class:`~sqlglot.expressions.Table`, when neither side of the + :class:`Intersects` predicate references the FROM table's alias, or + when the join's ``kind`` is outside the supported set + (``INNER`` / ``CROSS`` / ``SEMI`` / ``ANTI``; ``CROSS`` folds to + ``INNER`` because per-chromosome partitioning makes them equivalent + inside the inner subquery). + """ + join_table = join.this + if not isinstance(join_table, exp.Table): + return None + from_alias = _normalize_alias(from_table.alias_or_name) + left_alias = _normalize_alias(intersects.this.table) + right_alias = _normalize_alias(intersects.expression.table) + if left_alias == from_alias: + l_user_alias = from_alias + l_table = from_table + r_user_alias = right_alias + r_table = join_table + elif right_alias == from_alias: + # Inner subquery template hardcodes "a" to the FROM side; swap so + # that contract holds regardless of which operand the user put + # first in INTERSECTS. + l_user_alias = from_alias + l_table = from_table + r_user_alias = left_alias + r_table = join_table + else: + return None + raw_kind = (join.args.get("kind") or "").upper() + if raw_kind in ("", "CROSS"): + kind: Literal["INNER", "SEMI", "ANTI"] = "INNER" + elif raw_kind in ("SEMI", "ANTI"): + kind = raw_kind # type: ignore[assignment] + else: + return None + return cls( + left_user_alias=l_user_alias, + right_user_alias=r_user_alias, + left_table=l_table, + right_table=r_table, + kind=kind, + ) + + +@dataclass(frozen=True) +class _CountOverlapsShape: + """The matched ``count_overlaps`` LEFT-join-with-COUNT shape (#209). + + Bundles the LEFT :class:`~sqlglot.expressions.Join`, the un-aliased left + (FROM) base table, the left/right user aliases (case-folded), the left-side + group-key SELECT items, and the single ``COUNT`` aggregate SELECT item. The + emission reuses the INNER path to build the per-chromosome IEJoin count and + wraps it in a zero-fill LEFT join against the distinct left keys. + """ + + the_join: exp.Join + left_table: exp.Table + left_alias: str + right_alias: str + key_items: tuple[exp.Expression, ...] + agg_item: exp.Alias + + +def _match_count_overlaps(query: exp.Expression) -> "_CountOverlapsShape | None": + """Return the count_overlaps shape when *query* matches it, else ``None``. + + Matches exactly ``SELECT , COUNT() AS FROM + a LEFT JOIN b ON GROUP BY + `` -- the only right-side reference is the COUNT argument, and the + projected left columns are the group keys. Every other shape (``COUNT(*)``, + non-COUNT or multiple aggregates, a WHERE / HAVING / ORDER BY / LIMIT, a + non-LEFT outer join, extra ON residuals) returns ``None`` so the caller keeps + the existing behaviour (the naive-predicate plan). + """ + if not isinstance(query, exp.Select): + return None + if query.args.get("with_") or query.args.get("distinct"): + return None + if any( + query.args.get(k) is not None for k in ("having", "order", "limit", "offset") + ): + return None + if query.args.get("group") is None or query.args.get("where") is not None: + return None + + joins = query.args.get("joins") or [] + if len(joins) != 1: + return None + the_join = joins[0] + if _normalize_alias(the_join.args.get("side") or "") != "left": + return None + if (the_join.args.get("kind") or "").upper() not in ("", "OUTER"): + return None + on = the_join.args.get("on") + if on is None: + return None + intersects = _find_column_intersects_in(on) + if intersects is None or _count_column_intersects(query) != 1: + return None + # v1 supports only a bare INTERSECTS ON (no residual join conditions). + if _strip_intersects(on, intersects) is not None: + return None + + from_node = query.args.get("from_") + from_table = from_node.this if from_node else None + join_table = the_join.this + if not isinstance(from_table, exp.Table) or not from_table.name: + return None + if not isinstance(join_table, exp.Table) or not join_table.name: + return None + left_alias = _normalize_alias(from_table.alias_or_name) + right_alias = _normalize_alias(join_table.alias_or_name) + if left_alias == right_alias: + return None + + key_items: list[exp.Expression] = [] + agg_item: exp.Alias | None = None + for sel in query.expressions: + target = sel.this if isinstance(sel, exp.Alias) else sel + while isinstance(target, exp.Paren): + target = target.this + if isinstance(target, exp.AggFunc): + # Exactly one COUNT() / COUNT(DISTINCT ) aggregate, + # aliased so the wrapper can reference and re-project it. + if agg_item is not None or not isinstance(target, exp.Count): + return None + if not isinstance(sel, exp.Alias) or not sel.output_name: + return None + arg = target.this + if isinstance(arg, exp.Distinct): + if len(arg.expressions) != 1: + return None + arg = arg.expressions[0] + if not ( + isinstance(arg, exp.Column) + and not isinstance(arg.this, exp.Star) + and _normalize_alias(arg.table) == right_alias + ): + return None + agg_item = sel + elif ( + isinstance(target, exp.Column) + and target.table + and not isinstance(target.this, exp.Star) + ): + if _normalize_alias(target.table) != left_alias: + return None + key_items.append(sel) + else: + return None + + if agg_item is None or not key_items: + return None + # The GROUP BY keys must be exactly the projected left columns. An extra + # group key produces more rows than the distinct projected keys (which the + # zero-fill base relation reconstructs and so cannot reproduce), and a + # projected key absent from the group would be an invalid aggregate query. + projected_key_exprs = { + (item.this if isinstance(item, exp.Alias) else item).sql(dialect="duckdb") + for item in key_items + } + group_exprs = {e.sql(dialect="duckdb") for e in query.args["group"].expressions} + if group_exprs != projected_key_exprs: + return None + # The zero-fill wrapper joins on the user-facing output names and re-projects + # the count under its alias, so those names must be distinct — otherwise a + # key column and the count alias collide and USING / COALESCE bind to the + # wrong column. + output_names = [item.output_name for item in key_items] + [agg_item.output_name] + if len(set(output_names)) != len(output_names): + return None + + return _CountOverlapsShape( + the_join=the_join, + left_table=from_table, + left_alias=left_alias, + right_alias=right_alias, + key_items=tuple(key_items), + agg_item=agg_item, + ) + + +class IntersectsDuckDBIEJoinTransformer: + """Transform column-to-column INTERSECTS joins into a DuckDB IEJoin pattern. + + Emits a two-step dynamic SQL output: + + 1. ``SET VARIABLE __giql_iejoin_ = COALESCE(, );`` + where ```` is a ``string_agg`` over the chromosome + partition that emits one per-chromosome overlap query joined by + ``UNION ALL``. For INNER, that is ``SELECT ... FROM (... WHERE chrom = + '') a JOIN (... WHERE chrom = '') b ON ``; for + SEMI / ANTI it is ``SELECT ... FROM (... WHERE chrom = '') a WHERE + [NOT] EXISTS (SELECT 1 FROM (... WHERE chrom = '') b WHERE + )``. + 2. ``SELECT FROM query(getvariable('__giql_iejoin_'))`` + + Each per-chromosome subquery uses pure inequality predicates which DuckDB + plans through its range-join family (``IE_JOIN`` or ``PIECEWISE_MERGE_JOIN``). + The dialect supports INNER, SEMI, and ANTI variants. Crucially, DuckDB only + selects ``IE_JOIN`` for INNER pure-inequality joins — a bare ``SEMI JOIN`` / + ``ANTI JOIN`` inequality join is planned as a ``BLOCKWISE_NL_JOIN`` (a nested + loop, quadratic). So the left-only shapes are emitted as a correlated + ``WHERE EXISTS`` (SEMI) / ``WHERE NOT EXISTS`` (ANTI) subquery instead of a + semi/anti join keyword, which reaches ``IE_JOIN`` while preserving exact + semantics (#208). SEMI / ANTI restrict the outer SELECT to left-side + projections only. + + See :doc:`/transpilation/performance` for the full join-shape support + matrix, fallback enumeration, and hard-error projection shapes. + + Notes: + + - Star projections (bare ``*`` and ``a.*`` / ``b.*``) decline to the + generic naive-predicate plan, which delegates star expansion to DuckDB + and projects every base-table column against the live schema. The + dialect cannot enumerate a star at transpile time — it only knows the + :class:`~giql.table.Table` config's genomic columns (chrom / start / + end / strand), not arbitrary user columns — so declining keeps the + projection identical across backends rather than silently narrowing + it (#202). + - A projection the rebuild cannot express but the naive plan handles — + an expression (``a.start + 1``), a window aggregate, a ``FILTER`` + clause, a scalar subquery, an aggregate nested in an expression + (``COUNT(*) * 2``), or a star nested in an aggregate argument + (``COUNT(a.*)`` / ``MIN(COLUMNS(*))``) — also declines to the naive + plan, keeping ``dialect="duckdb"`` consistent with every backend + rather than hard-erroring or miscompiling (#204, #205). Only a + genuine out-of-scope column reference (unqualified / unknown-table / + right-side) still raises. + - ``SEMI JOIN`` and ``ANTI JOIN`` outputs are left-side only. Any + ``b.col`` reference in the outer SELECT, aggregate arguments, GROUP BY, + HAVING, or ORDER BY raises :class:`_UnqualifiedProjectionError` + (wrapped to :class:`ValueError` at the public boundary). A right-side + ``b.*`` under SEMI / ANTI declines to the naive plan with every other + star (which then rejects the out-of-scope right table at bind time). + - ``ANTI JOIN`` partitions on the left table's distinct chromosomes + only (not the chromosome INTERSECT used by INNER / SEMI) so that + left rows on chromosomes absent from the right table are preserved. + """ + + def __init__(self, tables: Tables): + """Initialize transformer. + + :param tables: + Table configurations for column mapping. + """ + self.tables = tables + + def transform_to_sql(self, query: exp.Expression) -> str | None: + """Return DuckDB IEJoin SQL for *query* if applicable, else None. + + The dialect rewrites the *whole* query as + ``SET VARIABLE __giql_iejoin_ = ...; SELECT ... FROM + query(getvariable(...))``, so any top-level clause :meth:`_build_sql` + does not read would be silently dropped. The check below is a + whitelist: engage the dialect path only for a query that is + exactly ``SELECT FROM + [JOIN ] {ON|WHERE} + `` with no other top-level clause. + Star projections (schema-less enumeration would narrow them, #202), + projections the rebuild cannot express but the naive plan handles — + expressions / window aggregates / FILTER clauses / scalar subqueries / + aggregate-nested stars (#204, #205) — and a ``SEMI`` / ``ANTI`` join + whose ``INTERSECTS`` sits out of scope in the ``WHERE`` (#201) also + decline. Everything else returns ``None`` so the generic + naive-predicate plan handles it. + """ + if not isinstance(query, exp.Select): + return None + + # Top-level WITH (CTEs) still falls back to the naive-predicate plan. + # ORDER BY / LIMIT / OFFSET / DISTINCT / GROUP BY / HAVING ride + # on the outer SELECT wrapper. ``DISTINCT ON (...)`` is the one + # DISTINCT variant we don't try to translate — it would interact + # with the outer alias rewrite in ways the naive-predicate plan handles + # for free. + if query.args.get("with_"): + return None + distinct_node = query.args.get("distinct") + if distinct_node is not None and distinct_node.args.get("on"): + return None + + # count_overlaps fast path (#209): a LEFT JOIN whose only right-side use + # is a COUNT(b.col) with a GROUP BY on the left keys reaches IE_JOIN via + # the INNER path plus a zero-fill LEFT join, rather than declining to the + # naive HASH_JOIN + inequality filter that every other outer join takes. + # Checked before the outer-join decline below since it *is* an outer join. + count_shape = _match_count_overlaps(query) + if count_shape is not None: + return self._build_count_overlaps_sql(query, count_shape) + + if _has_outer_join_intersects(query): + return None + + # A SEMI / ANTI join whose INTERSECTS lives in the top-level WHERE + # (out of scope for the right table) diverges from the reference + # plans, which reject it with a binder error. Decline so the naive + # plan surfaces that error instead of inventing results (#201). + if _has_left_only_join_where_intersects(query): + return None + + # Star projections (bare ``*`` or ``a.*`` / ``b.*``) need column + # enumeration the schema-less transpile cannot do without silently + # narrowing to the configured genomic columns. Decline so the naive + # plan expands the real star against DuckDB's live schema (#202). + if _has_star_projection(query): + return None + + if _count_column_intersects(query) != 1: + return None + + # The wrapper-relation rewriter is not scope-aware: any subquery in + # GROUP BY / HAVING / ORDER BY (or inside a SELECT-list aggregate + # argument) would have its inner column refs substituted with + # wrapper-relation aliases that don't exist in the subquery's own + # scope. Reject those shapes here. ``EXISTS (SELECT ...)`` parses + # as ``Exists(Select(...))`` without an ``exp.Subquery`` wrapper, + # so we also check ``exp.Select`` directly. A bare scalar subquery + # in the SELECT list (``SELECT (SELECT SUM(x) FROM ...)``) is not + # gated here — it declines to the naive plan later via the + # :meth:`_resolve_projections` projection pre-scan (#205), since its + # columns resolve against its own scope. + for clause_key in ("group", "having", "order"): + modifier = query.args.get(clause_key) + if modifier is None: + continue + if ( + modifier.find(exp.Subquery) is not None + or modifier.find(exp.Select) is not None + ): + return None + for sel in query.expressions: + target = sel.this if isinstance(sel, exp.Alias) else sel + while isinstance(target, exp.Paren): + target = target.this + if isinstance(target, exp.AggFunc) and ( + target.find(exp.Subquery) is not None + or target.find(exp.Select) is not None + ): + return None + + # The supported shape connects exactly two tables via one INTERSECTS, + # so the top level has exactly one join (either an explicit JOIN ... + # ON or an implicit comma-join in WHERE). A third comma/CROSS-joined + # table would otherwise be silently dropped by _build_sql. + joins = query.args.get("joins") or [] + if len(joins) != 1: + return None + the_join = joins[0] + + # NATURAL needs schema introspection we don't have at transpile time + # (Table configs only expose chrom/start/end/strand; user columns + # like name/score are invisible). The naive-predicate plan defers to + # DuckDB, which does see the schema, so falling back IS the optimal plan. + if the_join.args.get("method"): + return None + + # USING admits only the single-column form whose column matches both + # tables' ``chrom_col`` (the per-chromosome partition is exactly the + # equi-join that ``USING(chrom)`` requests). The chrom-equivalence + # check itself happens in ``_build_sql`` after ``l_table`` / + # ``r_table`` are resolved. Multi-column USING and USING(non-chrom) + # fall back; inline support for those is a documented follow-up. + using_cols = the_join.args.get("using") or [] + if using_cols and len(using_cols) != 1: + return None + + # INNER (the default), CROSS (functionally equivalent to INNER once + # the chromosome partition is in place), SEMI, and ANTI all engage. + # Everything else falls back so the naive-predicate plan can preserve its + # semantics (e.g. RIGHT / FULL outer joins, which the naive overlap + # predicate expresses as a plain outer ``ON`` condition). + kind_str = the_join.args.get("kind") + if kind_str and kind_str.upper() not in ("INNER", "CROSS", "SEMI", "ANTI"): + return None + + intersects, the_join = self._find_target_join(query) + if intersects is None or the_join is None: + return None + + from_table = query.args["from_"].this + if not isinstance(from_table, exp.Table): + return None + + # Reject non-base-table operands. A GIQL table-function operand such + # as ``DISJOIN(genes)`` parses as an ``exp.Table`` with an empty + # ``name`` and would be interpolated as an empty identifier into + # broken SQL. (A CTE-named operand is already handled by the + # ``with_`` guard above; an unregistered base table is fine — the + # dialect uses default columns just like the naive-predicate path does.) + if not from_table.name: + return None + join_table = the_join.this + if not isinstance(join_table, exp.Table) or not join_table.name: + return None + + sides = _IEJoinSides.from_intersects(from_table, the_join, intersects) + if sides is None: + return None + + if sides.left_user_alias == sides.right_user_alias: + # Same alias on both sides of the INTERSECTS — the inner subquery + # would alias the same underlying relation as both "a" and "b", + # which the naive-predicate plan renders correctly in place (no + # inner subquery, so no ambiguous double-alias). + return None + + # Compare fully-qualified identifiers (catalog/db/name) so two + # distinct same-named tables in different schemas are not treated + # as a self-join. Same qualified identifier still falls back so the + # naive-predicate plan handles the legitimate self-join shape. + if self._qualified_table_ident(sides.left_table) == self._qualified_table_ident( + sides.right_table + ): + return None + + # Extra predicates are inlined into each per-chromosome subquery's + # join ON, except WHERE residuals under a left-only (SEMI / ANTI) join, + # which :meth:`_build_sql` applies as an outer wrapper filter (#200). + # Soft-fallback residuals (subquery, + # aggregate, window, or an INTERSECTS still wrapped in OR/NOT/Paren) + # cause ``_build_sql`` to return None. User-mistake residuals + # (unqualified or unknown-aliased columns, or columns with a + # catalog/schema qualifier) raise :class:`_UnqualifiedProjectionError`, + # which we wrap to :class:`ValueError` here so the public surface is + # uniform. + try: + return self._build_sql(query, intersects, sides, the_join) + except _DeclineIEJoin: + # A projection the naive plan handles but the IEJoin cannot rebuild + # (expression / window / FILTER / scalar subquery / aggregate-nested + # star) — fall back to the naive-predicate plan (#204, #205). + return None + except _UnqualifiedProjectionError as exc: + raise ValueError(str(exc)) from exc + + def _build_count_overlaps_sql( + self, query: exp.Select, shape: "_CountOverlapsShape" + ) -> str | None: + """Render the count_overlaps zero-fill SQL for a matched LEFT-join shape (#209). + + Reuses the INNER IEJoin path: an INNER copy of the query (same + projections and GROUP BY, LEFT downgraded to INNER) is transpiled by + :meth:`transform_to_sql` into ``SET VARIABLE ...; ``, which + DuckDB plans through ``IE_JOIN`` + hash aggregate. The INNER count drops + left rows with no overlap, so its ``SELECT`` is wrapped as a CTE and + LEFT-joined back onto the distinct left keys, zero-filling the missing + counts. Returns ``None`` (fall back to the naive plan) if the INNER copy + itself declines. + """ + inner_query = query.copy() + inner_join = inner_query.args["joins"][0] + inner_join.set("side", None) + inner_join.set("kind", None) + inner_sql = self.transform_to_sql(inner_query) + if inner_sql is None or ";\n" not in inner_sql: + return None + set_var_stmt, inner_select = inner_sql.split(";\n", 1) + + q = self._sql_quote_ident + key_names = [item.output_name for item in shape.key_items] + key_projections = [item.sql(dialect="duckdb") for item in shape.key_items] + group_exprs = [ + (item.this if isinstance(item, exp.Alias) else item).sql(dialect="duckdb") + for item in shape.key_items + ] + agg_name = shape.agg_item.output_name + + # Distinct left-key relation, rendering the FROM table with its own alias + # so the key projections (``a.chrom`` ...) resolve against it. + base_cte = ( + f"SELECT {', '.join(key_projections)} " + f"FROM {shape.left_table.sql(dialect='duckdb')} " + f"GROUP BY {', '.join(group_exprs)}" + ) + using_cols = ", ".join(q(name) for name in key_names) + base_out = ", ".join(f"base.{q(name)}" for name in key_names) + wrapper = ( + f"WITH __giql_counts AS ({inner_select}), " + f"__giql_base AS ({base_cte}) " + f"SELECT {base_out}, COALESCE(c.{q(agg_name)}, 0) AS {q(agg_name)} " + f"FROM __giql_base AS base " + f"LEFT JOIN __giql_counts AS c USING ({using_cols})" + ) + return set_var_stmt + ";\n" + wrapper + + @staticmethod + def _sql_escape(s: str) -> str: + """Return *s* with single quotes doubled for safe SQL string literal use.""" + return s.replace("'", "''") + + @staticmethod + def _sql_quote_ident(name: str) -> str: + """Render *name* as a DuckDB-quoted identifier. + + Delegates to sqlglot's identifier emitter so embedded double quotes + are doubled correctly and DuckDB-specific quoting rules (case, + reserved words) are honored. Used at every identifier-interpolation + site in the dialect's dynamic SQL. + """ + return exp.to_identifier(name, quoted=True).sql(dialect="duckdb") + + @staticmethod + def _qualified_table_ident(table: exp.Table) -> str: + """Render *table*'s catalog/db/name as a qualified DuckDB identifier. + + Drops the user-supplied alias (``peaks AS a``) so the rendered string + is suitable both as a bare ``FROM`` operand (where the caller adds + its own alias) and as a subquery source. Catalog and schema qualifiers + survive — sqlglot's :meth:`~sqlglot.expressions.Expression.sql` emitter + handles identifier quoting per DuckDB's rules (reserved words and + special characters get quoted; safe unquoted identifiers stay + unquoted). + """ + no_alias = table.copy() + no_alias.set("alias", None) + return no_alias.sql(dialect="duckdb") + + @staticmethod + def _classify_extras( + extras: list[exp.Expression], + ) -> Literal["inline", "fallback"]: + """Return ``"fallback"`` if any extra needs the naive-predicate plan. + + Soft-fallback shapes: an INTERSECTS still embedded in the residual + (because :func:`_strip_intersects` couldn't peel an OR/NOT/Paren + wrapper), a subquery, an aggregate function, or a window function. + Returns ``"inline"`` when every extra is safe to inline into the + per-chromosome subquery's join ON clause. + """ + for predicate in extras: + if predicate.find(Intersects) is not None: + return "fallback" # INTERSECTS wrapped in OR/NOT/Paren + if predicate.find(exp.Subquery) is not None: + return "fallback" + if predicate.find(exp.Select) is not None: + return "fallback" + if predicate.find(exp.AggFunc) is not None: + return "fallback" + if predicate.find(exp.Window) is not None: + return "fallback" + return "inline" + + @staticmethod + def _validate_extra_qualifiers( + extras: list[exp.Expression], + sides: "_IEJoinSides", + ) -> None: + """Raise for unqualified, unknown-aliased, or schema-qualified columns. + + Surfaces user mistakes at transpile time rather than deferring to a + downstream binder error. Three failure modes raise: + + - **Unqualified column:** no ``.table`` qualifier. + - **Unknown alias:** ``.table`` doesn't match either side of the join. + - **Catalog/schema-qualified column:** ``mycat.myschema.a.score`` + would silently mis-bind once the alias-rewriter remaps ``a`` to + the inner subquery's hardcoded ``a`` alias (the ``mycat.myschema`` + qualifier survives the rewrite and addresses a different relation + than intended). + """ + valid_aliases = {sides.left_user_alias, sides.right_user_alias} + for predicate in extras: + for col in predicate.find_all(exp.Column): + col_table = _normalize_alias(col.table) if col.table else "" + if not col_table: + raise _UnqualifiedProjectionError( + f"dialect='duckdb' cannot inline extra predicate " + f"{predicate.sql()!r}: column {col.sql()!r} must be " + f"qualified with {sides.left_user_alias!r} or " + f"{sides.right_user_alias!r}." + ) + if col_table not in valid_aliases: + raise _UnqualifiedProjectionError( + f"dialect='duckdb' cannot inline extra predicate " + f"{predicate.sql()!r}: column references unknown " + f"table qualifier {col.table!r}; expected " + f"{sides.left_user_alias!r} or " + f"{sides.right_user_alias!r}." + ) + if col.args.get("db") or col.args.get("catalog"): + raise _UnqualifiedProjectionError( + f"dialect='duckdb' cannot inline extra predicate " + f"{predicate.sql()!r}: column {col.sql()!r} carries " + f"a catalog/schema qualifier; use the alias-only " + f"form ({sides.left_user_alias!r}. or " + f"{sides.right_user_alias!r}.)." + ) + + @staticmethod + def _rewrite_refs_for_per_chrom_subquery( + node: exp.Expression, + sides: "_IEJoinSides", + ) -> str: + """Render a residual predicate against the inner subquery's ``a``/``b`` scope. + + The user's AST references their original aliases (``peaks a`` / + ``genes b`` or ``peaks p`` / ``genes g``), but the per-chromosome + inner subquery is emitted with hardcoded aliases ``a`` and ``b`` + regardless of what the user wrote. Rewrite each column reference + accordingly before emitting the residual SQL into the inner join's + ON clause. Alias matching is case-folded so mixed-case identifiers + match the DuckDB-equivalent contract. + + Assumes residuals have already been validated to reference only + ``sides.left_user_alias`` or ``sides.right_user_alias``; columns + without a table qualifier or pointing at an unknown alias must + have been rejected upstream. + """ + + def replace(n: exp.Expression) -> exp.Expression: + if not isinstance(n, exp.Column): + return n + col_table = _normalize_alias(n.table) if n.table else "" + if col_table == sides.left_user_alias: + new = n.copy() + new.set("table", exp.to_identifier(sides.left_inner_alias)) + return new + if col_table == sides.right_user_alias: + new = n.copy() + new.set("table", exp.to_identifier(sides.right_inner_alias)) + return new + return n + + rewritten = node.transform(replace, copy=True) + return rewritten.sql(dialect="duckdb") + + @staticmethod + def _rewrite_refs_for_wrapper_relation( + node: exp.Expression, + projection_map: dict[tuple[str, str], str], + sides: "_IEJoinSides", + clause_label: str, + ) -> str: + """Rewrite table-qualified column refs to the wrapper relation's inner alias. + + The dialect's outer ``SELECT`` reads from ``query(getvariable(...)) AS + ``. The wrapper relation's columns are whatever the inner + per-chromosome subqueries projected — i.e., the ``__giql_p`` names + on the left-hand side of each ``AS`` in the inner SELECT. Any user + ``ORDER BY`` / ``GROUP BY`` / ``HAVING`` (and any aggregate-function + argument) written against the original table aliases (``a.start``, + ``b.score``) must be rewritten to those inner alias names so DuckDB + can resolve them against the wrapper relation. Alias matching is + case-folded for DuckDB-equivalent semantics. + + An upstream pre-allocation step ensures every referenced column has + an entry in ``projection_map``; missing keys raise + :class:`_UnqualifiedProjectionError` as a guard against future scope + relaxations that would otherwise produce silent miscompiles. + + Scope caveat: this walker is not scope-aware — it rewrites every + :class:`~sqlglot.expressions.Column` whose ``.table`` matches either + side's alias anywhere inside ``node``, including correlated or + non-correlated subqueries that may rebind those aliases to a + different table. The dispatcher gates already reject any subquery + in GROUP BY / HAVING / ORDER BY / SELECT list so the unsafe shape + can't reach here today. A future relaxation of those gates must + add a scope-aware walker here. + """ + valid_aliases = (sides.left_user_alias, sides.right_user_alias) + + def replace(n: exp.Expression) -> exp.Expression: + if not isinstance(n, exp.Column): + return n + col_table = _normalize_alias(n.table) if n.table else "" + if col_table not in valid_aliases: + return n + key = (col_table, n.name) + if key not in projection_map: + raise _UnqualifiedProjectionError( + f"dialect='duckdb' cannot resolve {n.sql()!r} in " + f"{clause_label}: ensure the column is referenced " + f"with a table qualifier matching the join's left " + f"or right alias." + ) + return exp.column(projection_map[key]) + + rewritten = node.transform(replace, copy=True) + return rewritten.sql(dialect="duckdb") + + def _extract_extra_predicates( + self, + query: exp.Select, + intersects: Intersects, + ) -> tuple[list[exp.Expression], list[exp.Expression]]: + """Return the non-INTERSECTS residuals split by provenance. + + Walks the AND tree under each JOIN ON and the WHERE, strips the + target INTERSECTS node, and returns ``(on_residuals, where_residuals)`` + in document order. ``on_residuals`` are the predicates that sat in a + ``JOIN ... ON`` alongside the INTERSECTS; ``where_residuals`` are the + predicates that sat in the top-level ``WHERE``. Both lists empty means + the query has no extras beyond the target INTERSECTS. + + The split matters because the two sources are *not* interchangeable + for every join kind. An ON residual is a genuine join condition and is + correct to inline into the per-chromosome ON for INNER / SEMI / ANTI + alike. A WHERE residual is a post-join filter: inlining it into the ON + is only sound for INNER (where ``ON overlap WHERE r`` ≡ ``ON overlap + AND r``) and SEMI (a left-only ``r`` is invariant over the right side); + for ANTI it inverts the anti-join for rows failing ``r`` (#200). See + :meth:`_build_sql`, which routes WHERE residuals to an outer filter for + the left-only (SEMI / ANTI) shapes. + + Note: :func:`_strip_intersects` only descends ``exp.And``. Predicates + whose AND tree wraps the INTERSECTS in ``exp.Or`` / ``exp.Not`` / + ``exp.Paren`` surface here with the INTERSECTS still embedded; + :meth:`_classify_extras` then routes them to the naive-predicate plan, while + :meth:`_validate_extra_qualifiers` enforces qualifier rules for the + remaining inlinable residuals. + """ + on_residuals: list[exp.Expression] = [] + for join in query.args.get("joins") or []: + on = join.args.get("on") + if on is None: + continue + residual = _strip_intersects(on, intersects) + if residual is not None: + on_residuals.append(residual) + where_residuals: list[exp.Expression] = [] + where = query.args.get("where") + if where: + residual = _strip_intersects(where.this, intersects) + if residual is not None: + where_residuals.append(residual) + return on_residuals, where_residuals + + def _find_target_join( + self, query: exp.Select + ) -> tuple[Intersects | None, exp.Join | None]: + """Locate the single INTERSECTS join target. + + Returns ``(intersects_node, join)`` where ``join`` is the + :class:`~sqlglot.expressions.Join` AST node carrying the target + INTERSECTS — either directly via its ``ON`` clause, or via an + implicit cross-join whose ``INTERSECTS`` predicate lives in the + top-level ``WHERE``. Alias matching on the WHERE-branch is + case-folded for DuckDB-equivalent semantics. + """ + for join in query.args.get("joins") or []: + on = join.args.get("on") + if on: + intersects = _find_column_intersects_in(on) + if intersects: + if not isinstance(join.this, exp.Table): + return None, None + return intersects, join + + where = query.args.get("where") + if where: + intersects = _find_column_intersects_in(where.this) + if intersects: + from_table = query.args["from_"].this + if not isinstance(from_table, exp.Table): + return None, None + from_alias = _normalize_alias(from_table.alias_or_name) + left_alias = _normalize_alias(intersects.this.table) + right_alias = _normalize_alias(intersects.expression.table) + target_alias = right_alias if left_alias == from_alias else left_alias + for join in query.args.get("joins") or []: + if isinstance(join.this, exp.Table): + if _normalize_alias(join.this.alias_or_name) == target_alias: + return intersects, join + + return None, None + + def _build_sql( + self, + query: exp.Select, + intersects: Intersects, + sides: "_IEJoinSides", + the_join: exp.Join, + ) -> str | None: + """Render the multi-statement DuckDB IEJoin SQL string for *query*. + + Returns ``None`` when a soft-fallback condition fires (the residual + of the join carries a shape the naive-predicate plan can handle but the + dialect cannot inline, or a single-column USING references a column + that is not both tables' ``chrom_col``). May propagate + :class:`_DeclineIEJoin` from the :meth:`_resolve_projections` projection + pre-scan (a naive-valid projection the rebuild cannot express — #204, + #205); :meth:`transform_to_sql` catches it and declines to the naive + plan. Raises :class:`_UnqualifiedProjectionError` when a user-mistake + condition fires; callers translating to the public surface wrap the + raise to :class:`ValueError`. + """ + on_residuals, where_residuals = self._extract_extra_predicates(query, intersects) + all_residuals = on_residuals + where_residuals + if all_residuals and self._classify_extras(all_residuals) == "fallback": + return None + if all_residuals: + self._validate_extra_qualifiers(all_residuals, sides) + + # ON residuals are genuine join conditions and inline into the + # per-chromosome ON for every kind. WHERE residuals are post-join + # filters: inlining them into the ON is only sound for INNER (and the + # equivalent SEMI case), but inverts the anti-join for rows failing the + # residual under ANTI (#200). For the left-only shapes (SEMI / ANTI) we + # therefore layer WHERE residuals as an outer filter on the wrapper + # relation instead of folding them into the join. INNER keeps inlining + # them (byte-identical output, provably equivalent). + if sides.is_left_only_join: + inline_residuals = on_residuals + outer_where_residuals = where_residuals + else: + inline_residuals = on_residuals + where_residuals + outer_where_residuals = [] + + # Tables registry is keyed by the bare table name; an unregistered + # table uses default column names like the naive-predicate plan does. + l_table = self.tables.get(sides.left_table.name) + r_table = self.tables.get(sides.right_table.name) + l_chrom = l_table.chrom_col if l_table else DEFAULT_CHROM_COL + l_start = l_table.start_col if l_table else DEFAULT_START_COL + l_end = l_table.end_col if l_table else DEFAULT_END_COL + r_chrom = r_table.chrom_col if r_table else DEFAULT_CHROM_COL + r_start = r_table.start_col if r_table else DEFAULT_START_COL + r_end = r_table.end_col if r_table else DEFAULT_END_COL + + # USING() admission check (the gate already restricted to + # single-column USING). The dialect's per-chromosome partition IS + # the equi-join that USING() requests; if the USING + # column doesn't match both tables' chrom_col (or the chrom_cols + # diverge), fall back so the naive-predicate plan preserves the + # USING equi-join for the engine to resolve. + using_cols = the_join.args.get("using") or [] + if using_cols: + using_name = _normalize_alias(using_cols[0].name) + if ( + _normalize_alias(l_chrom) != _normalize_alias(r_chrom) + or _normalize_alias(l_chrom) != using_name + ): + return None + + ( + inner_projections, + outer_projections, + projection_map, + ) = self._resolve_projections( + query, + sides, + outer_where_residuals, + ) + + var_name = new_variable_name("__giql_iejoin") + + # Canonicalize endpoints to 0-based half-open and use strict + # operators. Mixing inclusive operators with raw closed-closed + # coordinates would silently report non-overlapping touching + # intervals as overlapping. + q = self._sql_quote_ident + l_start_expr = canonical_start(f"a.{q(l_start)}", l_table) + l_end_expr = canonical_end(f"a.{q(l_end)}", l_table) + r_start_expr = canonical_start(f"b.{q(r_start)}", r_table) + r_end_expr = canonical_end(f"b.{q(r_end)}", r_table) + + # Render each table operand qualified (catalog/db/name) without the + # user's alias so the empty-schema fallback can supply its own + # ``a``/``b`` alias without producing ``FROM peaks AS a a`` shapes. + l_table_ident = self._qualified_table_ident(sides.left_table) + r_table_ident = self._qualified_table_ident(sides.right_table) + + # The overlap predicate (canonical, strict) plus any inline residuals, + # shared by both the INNER join-ON form and the SEMI / ANTI EXISTS form. + overlap_predicates = [ + f"{l_start_expr} < {r_end_expr}", + f"{l_end_expr} > {r_start_expr}", + ] + for residual in inline_residuals: + overlap_predicates.append( + self._rewrite_refs_for_per_chrom_subquery(residual, sides) + ) + predicate_sql = " AND ".join(overlap_predicates) + + # The dynamic SQL builder, expressed as a SQL string-concat that + # ``string_agg`` aggregates per-chromosome. Single quotes are + # doubled because the result is itself an SQL string literal that + # contains SQL. The chromosome literal is interpolated twice — once + # into the left partition filter, once into the right. + chrom_literal = CHROM_LITERAL + esc = self._sql_escape + inner_select_list = ", ".join(inner_projections) + + if sides.is_left_only_join: + # SEMI / ANTI: a bare per-chromosome ``SEMI JOIN`` / ``ANTI JOIN`` + # with an inequality overlap predicate is planned by DuckDB as a + # ``BLOCKWISE_NL_JOIN`` (a nested loop), not its fast ``IE_JOIN`` + # range join — DuckDB selects IE_JOIN only for INNER pure-inequality + # joins (#208). A correlated ``WHERE EXISTS`` / ``WHERE NOT EXISTS`` + # over the same per-chromosome right partition reaches IE_JOIN and + # preserves exact SEMI / ANTI semantics (one row per qualifying left + # row, no dedup-on-projection hazard). Inline (ON) residuals sit in + # the subquery WHERE beside the overlap predicate; WHERE residuals + # are still layered as an outer wrapper filter (``outer_where_residuals`` + # below, #200). + exists_keyword = "EXISTS" if sides.kind == "SEMI" else "NOT EXISTS" + select_from_left = ( + f"SELECT {inner_select_list} " + f"FROM (SELECT * FROM {l_table_ident} WHERE {q(l_chrom)} = " + ) + exists_open = ( + f") a WHERE {exists_keyword} (SELECT 1 FROM " + f"(SELECT * FROM {r_table_ident} WHERE {q(r_chrom)} = " + ) + exists_close = f") b WHERE {predicate_sql})" + per_chrom_sql_expr = ( + f"'{esc(select_from_left)}' " + f"|| {chrom_literal} " + f"|| '{esc(exists_open)}' " + f"|| {chrom_literal} " + f"|| '{esc(exists_close)}'" + ) + else: + # INNER (CROSS folded to INNER inside ``_IEJoinSides``): a + # per-chromosome INNER join whose pure-inequality ON predicate DuckDB + # plans through ``IE_JOIN``. + per_chrom_select_prefix = f"SELECT {inner_select_list} " + from_left = f"FROM (SELECT * FROM {l_table_ident} WHERE {q(l_chrom)} = " + from_right_open = ( + f") a JOIN (SELECT * FROM {r_table_ident} WHERE {q(r_chrom)} = " + ) + on_clause = f") b ON {predicate_sql}" + per_chrom_sql_expr = ( + f"'{esc(per_chrom_select_prefix)}{esc(from_left)}' " + f"|| {chrom_literal} " + f"|| '{esc(from_right_open)}' " + f"|| {chrom_literal} " + f"|| '{esc(on_clause)}'" + ) + + # Empty-schema fallback when no chromosomes pass the partition + # filter. Project from the source tables under WHERE FALSE so + # DuckDB resolves every output column to its real declared type. + # SEMI / ANTI outputs reference only the left side, so we drop + # the right operand from the empty-schema FROM clause. + empty_select_list = ", ".join(inner_projections) + if sides.is_left_only_join: + empty_schema = ( + f"SELECT {empty_select_list} FROM {l_table_ident} a WHERE FALSE" + ) + else: + empty_schema = ( + f"SELECT {empty_select_list} " + f"FROM {l_table_ident} a, {r_table_ident} b WHERE FALSE" + ) + + # Partition source: INTERSECT for INNER / SEMI (a chromosome only + # on one side produces no matches anyway), left-distinct for ANTI + # (a chromosome present only on the left must still emit its rows + # under ANTI semantics). + if sides.kind == "ANTI": + chrom_partition_subquery = ( + f"SELECT DISTINCT {q(l_chrom)} AS chrom FROM {l_table_ident}" + ) + else: + chrom_partition_subquery = ( + f"SELECT DISTINCT {q(l_chrom)} AS chrom FROM {l_table_ident} " + f"INTERSECT SELECT DISTINCT {q(r_chrom)} AS chrom " + f"FROM {r_table_ident}" + ) + + set_var_stmt = set_variable_statement( + var_name, per_chrom_sql_expr, chrom_partition_subquery, empty_schema + ) + + # Constant wrapper alias — uniqueness is provided by the + # token-bearing session variable; the outer SELECT's relation + # alias is never user-visible and doesn't need to vary per call. + wrapper_alias = "__giql_iejoin_wrapper" + outer_projection = ", ".join(outer_projections) + distinct_kw = "DISTINCT " if query.args.get("distinct") else "" + outer_select_parts = [ + f"SELECT {distinct_kw}{outer_projection} " + f"FROM {dynamic_relation(var_name, wrapper_alias)}" + ] + + # WHERE residuals for the left-only (SEMI / ANTI) shapes are applied + # here as a post-join filter on the wrapper relation, rewritten to the + # wrapper's inner-alias column names — never folded into the per-chrom + # ON, which would invert the anti-join for rows failing the residual + # (#200). Every referenced column was pre-allocated into the wrapper's + # projection list by :meth:`_resolve_projections`. + if outer_where_residuals: + where_sql = " AND ".join( + self._rewrite_refs_for_wrapper_relation( + residual, projection_map, sides, "WHERE" + ) + for residual in outer_where_residuals + ) + outer_select_parts.append(f"WHERE {where_sql}") + + # The rewriter translates any ``a.col`` / ``b.col`` references in + # the GROUP BY / HAVING / ORDER BY clauses to the wrapper + # relation's inner-alias column names (``__giql_p``), since + # the user's original aliases are not in scope in the outer SELECT. + group_node = query.args.get("group") + if group_node is not None: + outer_select_parts.append( + self._rewrite_refs_for_wrapper_relation( + group_node, projection_map, sides, "GROUP BY" + ) + ) + + having_node = query.args.get("having") + if having_node is not None: + outer_select_parts.append( + self._rewrite_refs_for_wrapper_relation( + having_node, projection_map, sides, "HAVING" + ) + ) + + order_node = query.args.get("order") + if order_node is not None: + outer_select_parts.append( + self._rewrite_refs_for_wrapper_relation( + order_node, projection_map, sides, "ORDER BY" + ) + ) + + limit_node = query.args.get("limit") + if limit_node is not None: + outer_select_parts.append(limit_node.sql(dialect="duckdb")) + + offset_node = query.args.get("offset") + if offset_node is not None: + outer_select_parts.append(offset_node.sql(dialect="duckdb")) + + outer_select = " ".join(outer_select_parts) + + return set_var_stmt + ";\n" + outer_select + + def _resolve_projections( + self, + query: exp.Select, + sides: "_IEJoinSides", + outer_where_residuals: list[exp.Expression] | None = None, + ) -> tuple[list[str], list[str], dict[tuple[str, str], str]]: + """Resolve the SELECT list (and any modifier-referenced columns). + + ``outer_where_residuals`` are the WHERE residuals that :meth:`_build_sql` + applies as an outer filter on the wrapper relation (the left-only + SEMI / ANTI shapes — #200). Their columns are pre-allocated into + ``inner_projections`` here so the wrapper exposes them even when they + appear only in the filter and not in the SELECT list. For INNER the + list is empty because those residuals inline into the per-chromosome ON. + + Returns ``(inner_projections, outer_projections, projection_map)``: + + * ``inner_projections`` is the per-chromosome subquery's SELECT + list — every column the query touches in any clause, projected + once each under a unique ``__giql_p`` alias. + * ``outer_projections`` is the outer SELECT's projection list, + rendered using the inner aliases (qualified columns and plain + aggregate calls). Star projections never reach here — they decline + upstream (#202) — nor do the naive-valid shapes the projection + pre-scan declines (expressions, window aggregates, FILTER clauses, + scalar subqueries, aggregate-nested stars — #204, #205). + * ``projection_map`` binds each ``(normalized_alias, column_name)`` + referenced anywhere in the query to its inner alias name so the + wrapper-relation rewriter and aggregate-argument rewriting can + translate user-written references to the wrapper relation. + + Raises :class:`_DeclineIEJoin` (caught upstream, declining to the naive + plan) for a SELECT-list shape the naive plan handles but the rebuild + cannot express. Raises :class:`_UnqualifiedProjectionError` only for a + genuine user error the naive plan also rejects — an unqualified column, + an unknown table qualifier, an aggregate over an unqualified column, or + (when ``sides.is_left_only_join``) a reference to the right side. + """ + q = self._sql_quote_ident + l_alias = sides.left_user_alias + r_alias = sides.right_user_alias + + # Decline (fall back to the naive-predicate plan) for any projection the + # IEJoin cannot rebuild but the naive plan handles — an expression, a + # window aggregate, a FILTER clause, a scalar subquery, an aggregate + # nested in an expression, or a star nested in an aggregate argument + # (#204, #205). Run before any allocation so an aggregate-nested star + # (COUNT(a.*)) declines rather than allocating a bogus ``a."*"`` column. + # A projection with an out-of-scope column reference is a user error the + # naive plan also rejects, so it is left to the dispatch below to raise. + for sel in query.expressions: + target = sel.this if isinstance(sel, exp.Alias) else sel + while isinstance(target, exp.Paren): + target = target.this + if _projection_declines_to_naive( + target, l_alias, r_alias, sides.is_left_only_join + ): + raise _DeclineIEJoin( + f"dialect='duckdb' cannot rebuild projection {target.sql()!r}; " + "deferring to the naive-predicate plan." + ) + + inner_projections: list[str] = [] + projection_map: dict[tuple[str, str], str] = {} + + def reject_right_side_if_left_only(tbl_normalized: str, source_sql: str) -> None: + """Raise when a SEMI / ANTI join projects from the right side.""" + if sides.is_left_only_join and tbl_normalized == r_alias: + raise _UnqualifiedProjectionError( + f"dialect='duckdb' with kind={sides.kind!r} only " + f"projects left-side columns; got {source_sql!r} " + f"which references the right side " + f"({sides.right_user_alias!r})." + ) + + def allocate(tbl: str, col: str) -> str: + """Return the inner alias for ``(tbl, col)``, allocating if new. + + Allocation order is deterministic per query but depends on + (a) the fixed pre-walk order in :meth:`_resolve_projections` — + modifier columns first (``("group", "having", "order")``), + then aggregate-argument columns, then the SELECT list — and + (b) sqlglot's ``find_all`` traversal order within each + sub-clause. Tests therefore should NOT couple to a specific + ``__giql_p`` ↔ source-column mapping (e.g. + ``assert projection_map[(a, start)] == "__giql_p3"``). + """ + tbl = _normalize_alias(tbl) + key = (tbl, col) + existing = projection_map.get(key) + if existing is not None: + return existing + if tbl == l_alias: + side = sides.left_inner_alias + elif tbl == r_alias: + side = sides.right_inner_alias + else: + raise _UnqualifiedProjectionError( + f"Unknown table qualifier {tbl!r} in projection; " + f"expected {l_alias!r} or {r_alias!r}." + ) + alias = f"__giql_p{len(inner_projections)}" + inner_projections.append(f"{side}.{q(col)} AS {alias}") + projection_map[key] = alias + return alias + + def render_aggregate(agg_node: exp.Expression, user_alias: str | None) -> None: + """Validate aggregate-argument qualifiers and emit the outer projection.""" + for col in agg_node.find_all(exp.Column): + col_table = _normalize_alias(col.table) if col.table else "" + if col_table not in (l_alias, r_alias): + raise _UnqualifiedProjectionError( + "dialect='duckdb' requires qualified aggregate " + f"arguments; got {agg_node.sql()!r}. Use a " + "qualified column reference like " + f"'{type(agg_node).__name__.upper()}(a.col)' or " + f"'{type(agg_node).__name__.upper()}(b.col)'." + ) + reject_right_side_if_left_only(col_table, agg_node.sql()) + rewritten_agg = self._rewrite_refs_for_wrapper_relation( + agg_node, projection_map, sides, "aggregate argument" + ) + outer_name = ( + user_alias + if user_alias is not None + else f"__giql_agg_{len(outer_projections)}" + ) + outer_projections.append(f"{rewritten_agg} AS {q(outer_name)}") + + # Pre-allocate inner aliases for every column referenced in + # GROUP BY / HAVING / ORDER BY so the wrapper relation exposes + # them even if the user didn't put them in the SELECT list. + for clause_key in ("group", "having", "order"): + modifier_node = query.args.get(clause_key) + if modifier_node is None: + continue + for col in modifier_node.find_all(exp.Column): + col_table = _normalize_alias(col.table) if col.table else "" + if col_table in (l_alias, r_alias): + reject_right_side_if_left_only(col_table, col.sql()) + allocate(col_table, col.name) + + # Pre-allocate inner aliases for columns referenced in the WHERE + # residuals routed to the outer wrapper filter (SEMI / ANTI, #200), so + # the wrapper exposes them even when they appear only in the filter. + # A right-side reference is rejected here (the left-only outer SELECT + # cannot resolve it), matching the SELECT-list / modifier behavior. + for residual in outer_where_residuals or []: + for col in residual.find_all(exp.Column): + col_table = _normalize_alias(col.table) if col.table else "" + if col_table in (l_alias, r_alias): + reject_right_side_if_left_only(col_table, col.sql()) + allocate(col_table, col.name) + + # Pre-allocate inner aliases for columns referenced inside any + # aggregate function in the SELECT list (we'll rewrite the + # aggregate's argument to use the inner alias when rendering the + # outer projection below). Peel any ``exp.Paren`` wrapper first + # so a paren-wrapped aggregate like ``(SUM(a.score)) AS s`` is + # treated identically to its bare form. + for sel in query.expressions: + target = sel.this if isinstance(sel, exp.Alias) else sel + while isinstance(target, exp.Paren): + target = target.this + if isinstance(target, exp.AggFunc): + for col in target.find_all(exp.Column): + col_table = _normalize_alias(col.table) if col.table else "" + if col_table in (l_alias, r_alias): + reject_right_side_if_left_only(col_table, col.sql()) + allocate(col_table, col.name) + + outer_projections: list[str] = [] + + for sel in query.expressions: + target = sel.this if isinstance(sel, exp.Alias) else sel + user_alias = sel.alias if isinstance(sel, exp.Alias) else None + # Peel paren wrappers at the top of dispatch so a paren-wrapped + # aggregate like ``(SUM(a.score)) AS s`` routes through the AggFunc + # branch below rather than the generic catch-all. + while isinstance(target, exp.Paren): + target = target.this + + # Star projections (bare ``*`` and ``a.*`` / ``b.*``) never reach + # here — :func:`_has_star_projection` declines them upstream (#202) — + # and neither do the naive-valid shapes the projection pre-scan + # above declined (expressions / window aggregates / FILTER clauses / + # scalar subqueries / aggregate-nested stars — #204, #205). + + if isinstance(target, exp.AggFunc): + render_aggregate(target, user_alias) + continue + + if isinstance(target, exp.Column) and target.table: + tbl = _normalize_alias(target.table) + reject_right_side_if_left_only(tbl, target.sql()) + col = target.name + alias = allocate(tbl, col) + outer_name = user_alias if user_alias is not None else col + outer_projections.append(f"{alias} AS {q(outer_name)}") + continue + + # Everything reaching here is an unsupported wrapper (an + # expression, a window aggregate, a FILTER clause, or an aggregate + # nested in an expression) that the pre-scan did NOT decline because + # it references an out-of-scope column — an unqualified column, an + # unknown table, or the right side under a SEMI / ANTI join (e.g. + # ``SUM(score) OVER (...)``, ``score + 1``). A plain aggregate with + # an unqualified argument (``SUM(score)``) is diagnosed earlier in + # ``render_aggregate``. + # + # A right-side column wrapped in an unsupported shape under a + # left-only join gets the dedicated left-side-only message rather + # than the generic "qualify the column" one, which would steer the + # user toward another invalid form (``b.col`` is never valid here). + for col in target.find_all(exp.Column): + col_table = _normalize_alias(col.table) if col.table else "" + reject_right_side_if_left_only(col_table, target.sql()) + # Otherwise the fault is an unqualified / unknown-table column; the + # accurate guidance is to qualify it (once qualified, the wrapper + # itself declines to the naive plan — #204, #205 — rather than + # erroring). + raise _UnqualifiedProjectionError( + "dialect='duckdb' requires qualified projections; got " + f"{target.sql()!r}. Use a qualified column reference like " + "'a.col' or 'b.col [AS x]'." + ) + + if not outer_projections: + raise _UnqualifiedProjectionError( + "dialect='duckdb' requires at least one qualified projection." + ) + + # An aggregate without any source-column argument (e.g. + # ``COUNT(*)``) leaves ``inner_projections`` empty in queries that + # have no other column references. The per-chromosome subquery + # still needs at least one column to project, so emit a constant + # placeholder that DuckDB optimizes away. (Without this the + # generated SQL ``SELECT FROM (...) a JOIN (...) b ON ...`` would + # fail to parse.) + if not inner_projections: + inner_projections.append("1 AS __giql_placeholder") + + return inner_projections, outer_projections, projection_map + + +def _has_sibling_spatial_predicate(node: Intersects, root: exp.Expression) -> bool: + """Return True if *root* holds a spatial predicate other than the join *node*. + + The former pre-pass ran on the raw parsed AST, so a residual GIQL spatial + predicate (``CONTAINS`` / ``WITHIN`` / ``ANY`` / ``ALL`` / a second + ``INTERSECTS``) sharing the query was always still an un-expanded GIQL node — + the transformer detected it and declined its IEJoin rewrite. In pass 3 the + ``ExpandOperators`` walk expands nodes deepest-first, so a sibling spatial + predicate may already have been rewritten to plain comparison SQL by the time + this override runs, which the transformer's ``_classify_extras`` can no longer + recognize (it would wrongly inline the residual into a per-chromosome ``ON``, + corrupting SEMI/ANTI results — #169). Detect the sibling either way: as an + un-expanded GIQL spatial node, or via the :data:`SPATIAL_PREDICATE_META` tag the + generic spatial expanders stamp on their output. Either signal means the join + must defer to the naive predicate, which composes correctly with any residual. + """ + spatial_types = (Intersects, Contains, Within, SpatialSetPredicate) + for candidate in root.walk(): + if candidate is node: + continue + if isinstance(candidate, spatial_types): + return True + if candidate.meta.get(SPATIAL_PREDICATE_META): + return True + return False + + +@register(DuckDBTarget, Intersects) +def expand_intersects_duckdb( + node: exp.Expression, ctx: ExpansionContext +) -> exp.Expression: + """Expand a DuckDB ``Intersects`` node — IEJoin join rewrite or naive predicate. + + The registered ``(DuckDBTarget, Intersects)`` override. For a column-to-column + INTERSECTS *join* whose whole-query shape the IEJoin path supports, + :meth:`IntersectsDuckDBIEJoinTransformer.transform_to_sql` builds the + per-chromosome ``SET VARIABLE ...; SELECT ... FROM query(getvariable(...))`` + string; since that is a whole-query multi-statement rewrite, not a node-local + expression, this registers a query-level finalizer + (:meth:`~giql.expander.ExpansionContext.add_statement_finalizer`) that replaces + the statement root with a :class:`~sqlglot.expressions.Command` wrapping the + built SQL (which serializes verbatim), and returns the ``Intersects`` node + unchanged (mirroring the CLUSTER / MERGE whole-query expanders). + + Every other shape — a join the IEJoin path declines (LEFT/RIGHT/FULL, self-join, + multiple INTERSECTS, extra predicates, non-base operands, 3+ tables), a + literal-range predicate, or a residual column-to-column predicate — defers to + :func:`giql.expanders.intersects._expand_spatial_op`, the same naive overlap + predicate the ``(GenericTarget, Intersects)`` expander emits on every target + (#167). A query carrying any *sibling* spatial predicate also defers (see + :func:`_has_sibling_spatial_predicate`), so the IEJoin never inlines a + residual it cannot see. ``transform_to_sql`` then runs on the pass-3 AST, which + for the shapes the IEJoin path accepts is structurally identical to the raw + parse — those shapes provably carry no other GIQL operator, and pass 1 only + attaches resolution metadata while pass 2 synthesizes no CTE for a + column-to-column INTERSECTS (its operands are column slots, not reference + slots) — so the rewrite reads the same query the former pre-pass did. + """ + if isinstance(node, Intersects) and _is_column_intersects(node): + root = node.root() + if isinstance(root, exp.Select) and not _has_sibling_spatial_predicate( + node, root + ): + transformer = IntersectsDuckDBIEJoinTransformer(ctx.tables) + iejoin_sql = transformer.transform_to_sql(root) + if iejoin_sql is not None: + ctx.add_statement_finalizer(lambda _root: exp.Command(this=iejoin_sql)) + return node + return _expand_spatial_op(node, ctx, "intersects") diff --git a/src/giql/expanders/merge.py b/src/giql/expanders/merge.py new file mode 100644 index 0000000..c9b3028 --- /dev/null +++ b/src/giql/expanders/merge.py @@ -0,0 +1,439 @@ +"""The MERGE operator expander (epic #137, issue #144). + +MERGE combines overlapping (and, with parameters, adjacent / strand-matched / +predicate-gated) intervals into single intervals. It is built on CLUSTER: assign +a cluster id, then aggregate ``MIN(start)`` / ``MAX(end)`` per cluster:: + + SELECT MERGE(interval) FROM features + +becomes:: + + SELECT chrom, MIN(start) AS start, MAX(end) AS end + FROM (SELECT *, CLUSTER(interval) AS __giql_cluster_id FROM features) + AS __giql_clustered + GROUP BY chrom, __giql_cluster_id + ORDER BY chrom, start + +(The ``becomes::`` form is simplified for readability; the emitted SQL quotes +identifiers, appends ``NULLS LAST``, and the inner ``CLUSTER(...)`` is itself +expanded into the two-level ``__giql_lag_calc`` form.) + +This module is the AST-expansion replacement for the legacy ``MergeTransformer`` +(in the former ``giql.transformer`` module, removed once every join rewrite became +registry-driven — #144, #169); it produces the same SQL (the existing MERGE +transpilation and bedtools oracle tests are the migration oracle). + +Like CLUSTER (:mod:`giql.expanders.cluster`), MERGE is a **whole-query rewrite**: +it navigates to the enclosing :class:`~sqlglot.expressions.Select`, mutates it in +place, and returns the operator node unchanged so the pass's ``node.replace`` is a +no-op. The intermediate clustered subquery is restructured through CLUSTER's shared +:func:`giql.expanders.cluster.expand_cluster_query`, mirroring how the legacy +``MergeTransformer`` composed ``ClusterTransformer``. +""" + +from __future__ import annotations + +from sqlglot import exp + +from giql.constants import CLUSTER_ID_COL +from giql.expander import ExpansionContext +from giql.expander import expand_operators +from giql.expander import register + +# The shared, operator-neutral expansion toolkit lives in the neutral ``_genomic`` +# module (#163). ``expand_cluster_query`` is the one deliberate CLUSTER-reuse +# point — MERGE composes the CLUSTER rewrite to build its clustered subquery — so +# it is imported from the CLUSTER expander as a documented composition rather than +# incidental coupling. +from giql.expanders._genomic import GenomicColumns +from giql.expanders._genomic import extract_stranded +from giql.expanders._genomic import find_projected +from giql.expanders._genomic import genomic_columns +from giql.expanders._genomic import is_star_projection +from giql.expanders._genomic import reject_cluster_merge_mix +from giql.expanders._genomic import require_top_level_projection +from giql.expanders._genomic import transplant +from giql.expanders.cluster import expand_cluster_query +from giql.expressions import GIQLCluster +from giql.expressions import GIQLMerge +from giql.targets import GenericTarget + + +@register(GenericTarget, GIQLMerge) +def expand_merge(node: GIQLMerge, ctx: ExpansionContext) -> exp.Expression: + """Expand a MERGE node by restructuring its enclosing SELECT in place. + + Registered for :class:`~giql.targets.GenericTarget` (MERGE emits identical SQL + across targets). Locates the enclosing + :class:`~sqlglot.expressions.Select`, derives the genomic columns from the FROM + table, and rewrites that SELECT in place into the clustered-aggregation form. + Returns the original MERGE node (now unreachable from the rewritten root SELECT) + so the pass's ``node.replace`` is a no-op. + + Parameters + ---------- + node : GIQLMerge + The MERGE node being expanded. + ctx : ExpansionContext + The expansion context; only its :attr:`~ExpansionContext.tables` is read. + + Returns + ------- + exp.Expression + The MERGE node, unchanged — the surrounding SELECT was mutated in place. + """ + select = node.find_ancestor(exp.Select) + if select is None: + return node + require_top_level_projection(select, node, GIQLMerge) + reject_cluster_merge_mix(select) + _reject_star_projection(select) + columns = genomic_columns(select, ctx.tables) + # Build from a detached copy, then transplant onto the original SELECT to + # preserve its identity (and rewrite a root SELECT without a parent to replace + # through), exactly as the CLUSTER expander does. + source = select.copy() + transformed = _transform_select_merge(source, columns) + if transformed is None: + # No-op for a MERGE that parses outside the SELECT projection (e.g. in + # WHERE / ORDER BY): find_projected finds none in the copy, so there is + # nothing to restructure and the operator leaks to the generator exactly + # as on `main`. require_top_level_projection above already rejects the + # in-projection-expression case; this guards the out-of-projection case. + return node + transplant(select, transformed) + # copy()+transplant duplicated the enclosing WHERE into the new __giql_clustered + # subquery; the originals the pass collected are now unreachable, so re-run the + # pass over the restructured SELECT to expand any sibling pass-3 operators + # carried into it. Safe from recursion: the MERGE is already gone. (#144 B1) + # + # expand_operators may return a new root (a registered statement finalizer can + # wrap it), and its contract requires callers to use the return value rather + # than assume in-place mutation, so reinstall a new root in place of `select`. + # Today the branch is never taken here — the sole finalizer-registering operator + # (a correlated NEAREST fallback) is already a plain join by this deepest-first + # re-walk, so this nested run registers no finalizer — but honoring the contract + # keeps the seam future-proof. A NEAREST fallback nested under this MERGE + # registered its finalizer in the outer expand_operators run; that finalizer + # re-locates the transplanted join by its reserved `meta` tag and would wrap a + # surfacing star (#172). MERGE's final projection is explicit, so it never + # surfaces the reserved columns — they are covered both by that (no-op) wrapper + # check and by the explicit projection. + result = expand_operators(select, ctx.target, ctx.tables, ctx.registry) + if result is not select: + select.replace(result) + return node + + +def _reject_star_projection(select: exp.Select) -> None: + """Raise if *select* projects a star alongside its MERGE. MERGE-specific guard. + + MERGE is an aggregate whole-query rewrite: its output is one row per merged + region (``chrom``, ``MIN(start)``, ``MAX(end)``), so a star projected in the + same SELECT has no coherent meaning — it names the *pre-aggregation* input + columns of a relation the rewrite has already collapsed into ``__giql_clustered`` + and grouped away. Left unguarded, a bare ``*`` re-surfaces those non-grouped, + non-aggregated columns under the synthesized ``GROUP BY`` and a qualified + ``rel.*`` dangles a table alias the outer aggregation no longer exposes — both + non-executable SQL (#189). Fail loudly here, mirroring the sibling MERGE guards + (multiple-MERGE, CLUSTER/MERGE mix), so the shape raises a clear diagnostic + instead. + + This is a genuinely MERGE-specific guard — unlike CLUSTER, a per-row window over + which a star is meaningful and supported (#184, #185), MERGE collapses rows, so + no star projection alongside it is well-formed — hence it lives here rather than + in the operator-neutral toolkit. The check is gated on there being a *projected* + MERGE (``find_projected``) so a MERGE parsing outside the projection (WHERE / + ORDER BY) is left on the expander's existing out-of-projection no-op path + rather than rejected. Legitimate sibling items — extra aggregates such as + ``COUNT(*)`` (a ``Star`` wrapped in an aggregate, not a bare or qualified star) + — are untouched; only a star is rejected. + """ + if not find_projected(select, GIQLMerge): + return + if any(is_star_projection(projection) for projection in select.expressions): + raise ValueError( + "MERGE cannot be combined with a star projection (e.g. SELECT *, " + "MERGE(...) or SELECT rel.*, MERGE(...)); MERGE aggregates rows into one " + "row per merged region, so a star has no coherent per-row meaning. Drop " + "the star, or project only grouping columns and aggregates (e.g. " + "COUNT(*)) alongside MERGE — not raw input columns, which are neither " + "grouped nor aggregated." + ) + + +def _transform_select_merge( + query: exp.Select, columns: GenomicColumns +) -> exp.Select | None: + """Rewrite *query* for the MERGE in its projection; ``None`` if none. + + Mirrors the legacy ``MergeTransformer.transform`` dispatch: find the MERGE + expressions, reject more than one, and delegate the single supported case to + :func:`_transform_for_merge`. + """ + merge_exprs = find_projected(query, GIQLMerge) + if not merge_exprs: + return None + # For now, support only one MERGE expression + if len(merge_exprs) > 1: + raise ValueError("Multiple MERGE expressions not yet supported") + return _transform_for_merge(query, merge_exprs[0], columns) + + +def _transform_for_merge( + query: exp.Select, merge_expr: GIQLMerge, columns: GenomicColumns +) -> exp.Select: + """Rewrite *query* into the clustered-aggregation form for one MERGE. + + A byte-for-byte port of the legacy ``MergeTransformer._transform_for_merge``, + with the genomic columns passed in and the intermediate clustered query + restructured through CLUSTER's shared + :func:`giql.expanders.cluster.expand_cluster_query` (the legacy method called + ``ClusterTransformer.transform``). Builds an inner ``__giql_clustered`` subquery + that appends ``__giql_cluster_id``, then an outer query that aggregates + ``MIN(start)`` / ``MAX(end)`` per cluster. + """ + chrom_col, start_col, end_col, strand_col = columns + + # Extract MERGE parameters (same as CLUSTER) + distance_expr = merge_expr.args.get("distance") + stranded_expr = merge_expr.args.get("stranded") + predicate_expr = merge_expr.args.get("predicate") + + # Build CLUSTER expression with same parameters + cluster_kwargs = {"this": merge_expr.this} + if distance_expr: + cluster_kwargs["distance"] = distance_expr + if stranded_expr: + cluster_kwargs["stranded"] = stranded_expr + if predicate_expr is not None: + cluster_kwargs["predicate"] = predicate_expr + + cluster_expr = GIQLCluster(**cluster_kwargs) + + # Create intermediate query with CLUSTER + # Start with original query's FROM/WHERE/etc + cluster_query = exp.Select() + cluster_query.select(exp.Star(), copy=False) + # __giql_cluster_id and its ``__giql_clustered`` / ``__giql_lag_calc`` / + # ``__giql_is_new_cluster`` siblings all carry the reserved prefix (#161). + cluster_query.select( + exp.alias_(cluster_expr, CLUSTER_ID_COL, quoted=False), + append=True, + copy=False, + ) + + # Copy FROM, WHERE from original + # Use copy() to avoid sharing references between queries + if query.args.get("from_"): + cluster_query.set("from_", query.args["from_"].copy()) + if query.args.get("where"): + cluster_query.set("where", query.args["where"].copy()) + + # Apply CLUSTER transformation to get the CTE-based query. hide_reserved=False: + # MERGE opts out of CLUSTER's flag-hiding (#184) because its explicit outer + # projection (chrom, MIN/MAX) never surfaces __giql_is_new_cluster, so hiding it + # would only add a needless ``* EXCEPT`` to this intermediate clustered subquery. + clustered = expand_cluster_query(cluster_query, columns, hide_reserved=False) + assert clustered is not None, "intermediate MERGE cluster query has no CLUSTER" + cluster_query = clustered + + # Build GROUP BY columns + # Quote the chrom column (as every other chrom reference is) so a reserved-word + # custom chrom column emits valid SQL (#144 A13). + group_by_cols = [exp.column(chrom_col, quoted=True)] + + # Handle stranded parameter + stranded = extract_stranded(stranded_expr) + + if stranded: + group_by_cols.append(exp.column(strand_col, quoted=True)) + + group_by_cols.append(exp.column(CLUSTER_ID_COL)) + + # The columns the synthesized grouping exposes to the user's projection and GROUP + # BY: chrom, plus strand when stranded (the per-merge __giql_cluster_id is internal + # and never user-referenceable). Used below to reconcile the user's own projection + # and GROUP BY against MERGE's aggregation (#192). + grouping_keys = {chrom_col} + if stranded: + grouping_keys.add(strand_col) + # Every identifier comparison below folds case. SQL binds unquoted identifiers + # case-insensitively, so ``AS Start`` collides with the synthesized ``start`` and + # ``GROUP BY CHROM`` names the ``chrom`` grouping key. Folding both sides matches a + # case-variant rather than silently missing it — the safe direction, since a missed + # collision is silently-wrong output while an over-match merely fails loud (#192). + grouping_keys_folded = {key.lower() for key in grouping_keys} + + # Reconcile the user's GROUP BY with MERGE's synthesized grouping. MERGE groups by + # chrom (and strand when stranded) plus its per-merge __giql_cluster_id, so a user + # GROUP BY over exactly those grouping-key columns is subsumed by it (the documented + # "merge by chromosome" shape). Anything else — another column, or an expression / + # ordinal MERGE cannot map to its grouping — was previously dropped silently, + # changing the result, so reject it loudly. Each item must be a *bare column* + # reference to a grouping key: an expression (GROUP BY UPPER(chrom)) or an ordinal + # (GROUP BY 1) is not an exp.Column, so it cannot be verified to name a grouping key + # and is rejected rather than honored-then-discarded (#192). + user_group = query.args.get("group") + if user_group is not None: + for item in user_group.expressions: + if ( + not isinstance(item, exp.Column) + or item.name.lower() not in grouping_keys_folded + ): + raise ValueError( + f"MERGE cannot honor GROUP BY {item.sql()}: MERGE groups rows by " + "chrom (and strand when stranded) and its synthesized per-merge " + "cluster id, so group only by a bare chrom (or strand) column " + "reference alongside MERGE — not another column, an expression, or " + "an ordinal, which it cannot faithfully honor." + ) + + # Build SELECT expressions for merged intervals + select_exprs = [] + + # Add group-by columns (non-aggregated) + select_exprs.append(exp.column(chrom_col, quoted=True)) + if stranded: + select_exprs.append(exp.column(strand_col, quoted=True)) + + # Add merged interval bounds + select_exprs.append( + exp.alias_( + exp.Min(this=exp.column(start_col, quoted=True)), start_col, quoted=False + ) + ) + select_exprs.append( + exp.alias_(exp.Max(this=exp.column(end_col, quoted=True)), end_col, quoted=False) + ) + + # The output-column names MERGE synthesizes: the grouping keys plus the aggregated + # ``start`` / ``end`` bounds. A user item may not silently collide with these. Folded + # for case-insensitive comparison, as grouping_keys_folded above. + synthesized_names_folded = { + name.lower() for name in (grouping_keys | {start_col, end_col}) + } + + # Process the other projection items. MERGE already projects the grouping keys and + # the MIN/MAX bounds and aggregates rows into one row per merged region, so each + # remaining item is reconciled against that grouping rather than copied verbatim — + # copying verbatim duplicated a grouping key the user also projected and, for a raw + # non-grouped column, emitted SQL the engine rejects (#192). + for expression in query.expressions: + # Skip the MERGE expression itself + if isinstance(expression, GIQLMerge): + continue + elif isinstance(expression, exp.Alias) and isinstance( + expression.this, GIQLMerge + ): + continue + + # A projected item is well-formed under MERGE's per-cluster GROUP BY only if + # every column it references *outside a grouping aggregate* is a grouping key. + # A column counts as ungrouped when no aggregate wraps it (a raw ``score``, or + # ``start`` in ``MAX(score) + start``) OR when the aggregate that wraps it is a + # *window* aggregate (``SUM(score) OVER (...)``): a window is evaluated after the + # GROUP BY and does not collapse the group, so its argument still has to be + # grouped. Fail loudly rather than emitting SQL the engine rejects (or that a + # lenient engine runs with an arbitrary-row value) (#192). + ungrouped = { + col.name.lower() + for col in expression.find_all(exp.Column) + if col.find_ancestor(exp.AggFunc) is None + or col.find_ancestor(exp.Window) is not None + } + if not ungrouped <= grouping_keys_folded: + raise ValueError( + "MERGE cannot project the non-aggregated column " + f"{', '.join(repr(c) for c in sorted(ungrouped - grouping_keys_folded))}" + f" in '{expression.sql()}': MERGE aggregates rows into one row per " + "merged region, so project only grouping columns (chrom, and " + "strand when stranded) and plain aggregates (e.g. COUNT(*)) " + "alongside MERGE." + ) + + out_name = expression.output_name + inner = expression.this if isinstance(expression, exp.Alias) else expression + + # A grouping-key column MERGE already projects, under its own name (bare + # ``chrom`` or a no-op ``chrom AS chrom`` of the "merge by chromosome" shape): + # drop the duplicate rather than surfacing the column twice (#192). + if ( + out_name.lower() in grouping_keys_folded + and isinstance(inner, exp.Column) + and inner.name.lower() == out_name.lower() + ): + continue + + # Reject an item whose *output name* collides with a synthesized column (the + # aggregated ``start`` / ``end`` bounds, or a value aliased onto a grouping key + # such as ``chrom AS start``): emitting two columns of that name silently + # duplicates one, and for ``start`` / ``end`` hijacks the default + # ``ORDER BY "chrom", "start"``, returning rows in the wrong order. Apply this + # only to an explicit alias or a bare column, whose output name the engine emits + # verbatim; for any other unaliased expression (``CAST(chrom AS VARCHAR)``, + # ``chrom || ''``), ``output_name`` is a sqlglot heuristic that does NOT match + # the engine's generated label, so checking it there would falsely reject a + # well-formed projection over a grouping key. Fold case so an unquoted ``Start`` + # is caught (it binds to the synthesized ``start``) (#192). + names_reliable = isinstance(expression, (exp.Alias, exp.Column)) + if names_reliable and out_name.lower() in synthesized_names_folded: + raise ValueError( + f"MERGE cannot project a column named {out_name!r} alongside MERGE: it " + "collides with the chrom/start/end columns MERGE synthesizes for each " + "merged region — alias it to a different name." + ) + + # An aggregate over the merged rows (COUNT(*), AVG(score), ...) or a + # non-aggregate expression over only grouping-key columns (UPPER(chrom), a + # rename ``chrom AS c``): keep it — well-formed under the GROUP BY, no collision. + select_exprs.append(expression) + + # Build final query + final_query = exp.Select() + final_query.select(*select_exprs, copy=False) + + # FROM the clustered subquery. The alias carries the reserved __giql_ prefix for + # namespace consistency with the other synthesized names (#161); a derived-table + # alias never actually collides with a user relation (it lives in the enclosing + # query's scope), so this is hygiene, not a fix. + subquery = exp.Subquery( + this=cluster_query, + alias=exp.TableAlias(this=exp.Identifier(this="__giql_clustered")), + ) + final_query.from_(subquery, copy=False) + + # Add GROUP BY + final_query.group_by(*group_by_cols, copy=False) + + # Honor the user's HAVING over the merged output. MERGE aggregates rows into one row + # per merged region, so a HAVING filters those regions — ``HAVING COUNT(*) > 1`` + # keeps only regions built from more than one input interval, resolving against the + # per-region GROUP BY above. The rewrite builds ``final_query`` fresh and transplant + # clears the original HAVING (a rebuilt root arg), so it must be re-attached here; + # omitting it silently dropped the clause and returned unfiltered rows (#192). An + # invalid HAVING (over a raw non-grouped column) surfaces a loud engine binder error, + # exactly as an unsupported ORDER BY does — never silently-wrong output. + user_having = query.args.get("having") + if user_having is not None: + final_query.set("having", user_having.copy()) + + # ORDER BY: honor the user's ORDER BY over the merged output when present, + # otherwise default to (chromosome, start) for deterministic output. MERGE + # aggregates rows away, so the user's ORDER BY resolves against the merged + # columns (e.g. ``ORDER BY "end" DESC`` over ``MAX("end") AS end``). Overriding + # it with the fixed default silently returned the wrong rows once an outer LIMIT + # was preserved (#181), so the user's ordering takes precedence. + if query.args.get("order"): + final_query.set("order", query.args["order"].copy()) + else: + final_query.order_by( + exp.Ordered(this=exp.column(chrom_col, quoted=True)), copy=False + ) + final_query.order_by( + exp.Ordered(this=exp.column(start_col, quoted=True)), + append=True, + copy=False, + ) + + return final_query diff --git a/src/giql/expanders/nearest.py b/src/giql/expanders/nearest.py new file mode 100644 index 0000000..46ff4c8 --- /dev/null +++ b/src/giql/expanders/nearest.py @@ -0,0 +1,826 @@ +"""The NEAREST operator expander (epic #137, issue #142). + +NEAREST is the first operator whose expansion is genuinely capability-driven. +The portable form is a correlated ``LATERAL`` subquery: each outer row drives a +``SELECT ... FROM WHERE ORDER BY ABS(distance) +LIMIT k`` whose reference endpoints are outer-table columns. DuckDB and the +generic target plan that directly (``supports_lateral == True``). + +Apache DataFusion has no correlated-``LATERAL`` physical plan +(``supports_lateral == False``). For it the same k-nearest / ``max_distance`` / +``stranded`` / ``signed`` semantics are reproduced with a **decorrelated +window-function fallback**: the target is cross-joined against the outer +relation, each candidate is ranked with +``ROW_NUMBER() OVER (PARTITION BY ORDER BY ABS(distance))``, and +the surrounding ``CROSS JOIN LATERAL`` is rewritten into a plain join that +re-associates the top-``k`` ranked candidates back to every outer row sharing +that reference key. Ranking depends only on the reference value, so ranking once +per distinct reference value and re-joining is set-equivalent to the per-row +LATERAL form — deterministic when the ``(start, end)`` tiebreaker distinguishes +candidates tied at the k-th distance (verified by the cross-target result +oracle). + +A literal-reference (standalone) NEAREST is already an uncorrelated subquery, so +every target — DataFusion included — uses the LATERAL/standalone form unchanged; +only the *correlated* shape needs the fallback. + +The expander assembles its distance/passthrough/filter SQL as string fragments +(the ``_nearest_*`` helpers below and +:func:`giql.expanders._distance.generate_distance_case`, shared in spirit with +DISTANCE, #140), then parses the assembled fragments into AST so the emitted SQL +is reserialized by the active target's serializer. +""" + +from __future__ import annotations + +from typing import Callable + +from sqlglot import exp +from sqlglot import parse_one + +from giql.canonical import decanonical_end +from giql.canonical import decanonical_start +from giql.dialect import GIQLDialect +from giql.expander import ExpansionContext +from giql.expander import StatementFinalizer +from giql.expander import register +from giql.expanders._distance import generate_distance_case +from giql.expanders._params import coerce_bool_param +from giql.expressions import GIQLNearest +from giql.range_parser import RangeParser +from giql.resolver import META_KEY +from giql.resolver import OperatorResolution +from giql.resolver import ResolvedInterval +from giql.resolver import ResolvedRef +from giql.table import Table +from giql.targets import Capabilities +from giql.targets import GenericTarget + +#: Meta key marking the decorrelated fallback's replacement subquery so its +#: statement finalizer can re-locate it in the live root at apply time (walking the +#: root, see :func:`_find_fallback_subquery`) rather than capturing the ``join`` +#: object eagerly (#172). The value is the fallback's unique per-run *tag* (see +#: :func:`_reserved_column_names`). A ``meta`` marker rides ``.copy()`` into the +#: tree a CLUSTER/MERGE ``copy()`` + ``transplant`` +#: (:func:`giql.expanders._genomic.transplant`) rebuilds, whereas a captured +#: reference would strand in the discarded original — the same reliance on ``meta`` +#: surviving ``.copy()`` that CLUSTER's re-entrant expansion depends on for the +#: resolution :data:`giql.resolver.META_KEY`. Derived from +#: :data:`giql.resolver.META_KEY` so it shares the single greppable ``giql`` +#: metadata namespace. +_FALLBACK_TAG_META = f"{META_KEY}_nearest_fallback_tag" + + +def _reserved_column_names(tag: str) -> tuple[list[str], str, str]: + """Return this fallback's reserved ``(key_names, rank_col, ref_key_prefix)``. + + The decorrelated window-function fallback synthesizes reserved rank/key columns + inside its ranked subquery. Their names are minted *per fallback* from *tag* — a + fresh :meth:`~giql.expander.ExpansionContext.alias` value (itself reserved by + the :data:`giql.expander.EXPAND_ALIAS_PREFIX` ``__giql_x_`` prefix) — rather than + fixed module constants, so two correlated NEAREST fallbacks in one query never + emit the same physical column name (#172). Identical names would collide in the + combined output (the engine disambiguates the second set with a ``…:1`` suffix that a + single ``* EXCEPT`` cannot name), leaking the second fallback's reserved + columns; distinct per-fallback names let each fallback's finalizer ``* EXCEPT`` + exactly its own set. The ``rk`` / ``rn`` suffixes and the reserved prefix keep + every name clear of user identifiers. + + :param tag: + The fallback's unique per-run tag (a minted alias, e.g. ``__giql_x_2``). + :return: + ``(key_names, rank_col, ref_key_prefix)`` — the ``[chrom, start, end]`` + reference-key column names, the ``ROW_NUMBER`` rank column name, and the + prefix from which the optional stranded ``strand`` key is derived. + """ + ref_key_prefix = f"{tag}_rk_" + key_names = [f"{ref_key_prefix}{part}" for part in ("chrom", "start", "end")] + rank_col = f"{tag}_rn" + return key_names, rank_col, ref_key_prefix + + +def _nearest_output_encoding( + expression: GIQLNearest, target_ref: ResolvedRef +) -> Table | None: + """Return the target's declared encoding for NEAREST's row passthrough. + + ``CanonicalizeCoordinates`` (pass 2) records the original + :class:`~giql.table.Table` on the resolution when it wraps a non-canonical + target in a ``__giql_canon_*`` CTE (blanking the slot's own ``table``). For + an unwrapped target — a canonical registered table, or any target when the + pass did not run — the slot's own ``table`` carries the (identity) encoding. + + :param expression: + GIQLNearest expression node + :param target_ref: + The resolved target reference (post pass 2) + :return: + The target's declared :class:`~giql.table.Table`, or ``None`` + """ + resolution = expression.meta.get(META_KEY) + if isinstance(resolution, OperatorResolution): + preserved = resolution.output_tables.get("this") + if preserved is not None: + return preserved + return target_ref.table + + +def _nearest_passthrough( + table_name: str, + target_start: str, + target_end: str, + output_table: Table | None, + capabilities: Capabilities, +) -> str: + """Project the target's full row, de-canonicalizing the interval columns. + + NEAREST passes the whole target row through (``SELECT {table_name}.*``) + alongside the synthesized, encoding-invariant ``distance`` column. When the + target's declared encoding is canonical 0-based half-open the row passes + through as a plain ``{table_name}.*`` — the byte-identical identity fast + path. When it is non-canonical the interval columns, canonical inside the + ``__giql_canon_*`` CTE the target was rewritten to, are de-canonicalized + back into that encoding so the passed-through interval matches the target's + own convention. + + The emit strategy is chosen from *capabilities*, following the precedent of + :func:`giql.expanders.disjoin._disjoin_passthrough` (issue #145) — the same + two emit forms: + + * ``{table_name}.* REPLACE (...)`` when ``supports_star_replace`` holds — + substitutes start/end in place; + * the portable ``{table_name}.* EXCEPT (start, end), , `` form + otherwise (the generic baseline / DataFusion family). Row-equivalent but + not column-order-equivalent, and not DuckDB-runnable. + + :param table_name: + The relation the row is selected from (the canon CTE name when wrapped, + else the registered table name) — also the column qualifier. + :param target_start: + Physical start column name + :param target_end: + Physical end column name + :param output_table: + The target's declared :class:`~giql.table.Table`, or ``None`` + :param capabilities: + The active target's :class:`~giql.targets.Capabilities`. + :return: + The passthrough projection fragment + """ + if output_table is None or ( + output_table.coordinate_system == "0based" + and output_table.interval_type == "half_open" + ): + return f"{table_name}.*" + pt_start = decanonical_start(f'{table_name}."{target_start}"', output_table) + pt_end = decanonical_end(f'{table_name}."{target_end}"', output_table) + if capabilities.supports_star_replace: + return ( + f"{table_name}.* REPLACE " + f'({pt_start} AS "{target_start}", {pt_end} AS "{target_end}")' + ) + # Portable form for engines without ``* REPLACE`` (generic / DataFusion): + # drop the interval columns from the star and re-project them recomputed. + return ( + f'{table_name}.* EXCEPT ("{target_start}", "{target_end}"), ' + f'{pt_start} AS "{target_start}", {pt_end} AS "{target_end}"' + ) + + +def _raise_nearest_reference_error( + expression: GIQLNearest, + resolution: OperatorResolution | None, +) -> None: + """Raise the historical diagnostic for an unresolved NEAREST reference. + + ResolveOperatorRefs (pass 1) defers a reference slot it cannot resolve; this + re-raises the historical pre-pass error verbatim. The implicit-outer failures + rely on the :class:`~giql.resolver.SlotDeferral` the pass records (the + ancestor walk that distinguished them now lives in the pass); a literal-range + parse failure is reproduced by re-parsing. + + :param expression: + GIQLNearest expression node + :param resolution: + The attached resolution metadata, if any + :raises ValueError: + Always — with the matching historical message + """ + reference = expression.args.get("reference") + + if reference is None: + # An absent reference is a correlated (implicit-outer) placement that the + # resolver could not tie to a registered outer table; consult the + # recorded deferral for the specific historical message. + deferral = resolution.deferral("reference") if resolution is not None else None + if deferral is not None and deferral.reason == "implicit_outer_unregistered": + raise ValueError( + f"Outer table '{deferral.detail}' not found in tables. " + "Please specify reference parameter explicitly." + ) + raise ValueError( + "Could not find outer table in LATERAL join context. " + "Please specify reference parameter explicitly." + ) + + # An explicit reference that deferred is a literal range that failed to + # parse (column references always resolve). Re-parse to surface the + # original parse error in the historical message. + reference_sql = reference.sql(dialect=GIQLDialect) + range_str = reference_sql.strip("'\"") + try: + RangeParser.parse(range_str).to_zero_based_half_open() + except Exception as e: + raise ValueError( + f"Could not parse reference genomic range: {range_str}. Error: {e}" + ) + raise ValueError(f"Could not parse reference genomic range: {range_str}.") + + +def _nearest_params( + expression: GIQLNearest, +) -> tuple[int, int | None, bool, bool]: + """Unpack the (k, max_distance, stranded, signed) parameters of a NEAREST.""" + k = expression.args.get("k") + k_value = int(str(k)) if k else 1 + + max_distance = expression.args.get("max_distance") + max_dist_value = int(str(max_distance)) if max_distance else None + + is_stranded = coerce_bool_param(expression.args.get("stranded")) + is_signed = coerce_bool_param(expression.args.get("signed")) + return k_value, max_dist_value, is_stranded, is_signed + + +def _distance_and_filters( + expression: GIQLNearest, + table_name: str, + target_ref: ResolvedRef, + ref: ResolvedInterval, + capabilities: Capabilities, + ref_fragments: tuple[str, str, str, str | None] | None = None, +) -> tuple[str, str, list[str], str]: + """Build the shared distance SQL, the qualified target columns, and WHERE. + + Returns ``(distance_expr, abs_distance_expr, where_clauses, passthrough)`` — + the fragments common to the LATERAL/standalone form and the decorrelated + fallback. Distance math, the chromosome pre-filter, the optional strand match, + and the optional ``max_distance`` filter all reproduce the legacy + ``giqlnearest_sql`` emitter exactly. Each form derives its deterministic + ORDER BY tiebreaker from the target columns itself. + + ``capabilities`` is the active target's :class:`~giql.targets.Capabilities`, + forwarded to :func:`_nearest_passthrough` to choose the target's + de-canonicalization emit form (``* REPLACE`` vs the portable ``* EXCEPT``); + both call sites pass ``ctx.capabilities``. + + ``ref_fragments`` optionally overrides the reference ``(chrom, start, end, + strand)`` SQL fragments. The LATERAL form consumes the resolution's + outer-qualified fragments verbatim; the fallback passes fragments pointing at + its renamed, pre-projected reference relation so the cross-joined columns + carry names distinct from the target's (DataFusion's planner cannot resolve a + window ordering over a join with duplicate column names). + """ + target_chrom, target_start, target_end = target_ref.cols + _k_value, max_dist_value, is_stranded, is_signed = _nearest_params(expression) + + output_table = _nearest_output_encoding(expression, target_ref) + passthrough = _nearest_passthrough( + table_name, target_start, target_end, output_table, capabilities + ) + + if ref_fragments is not None: + ref_chrom, ref_start, ref_end, ref_strand_frag = ref_fragments + else: + ref_chrom, ref_start, ref_end, ref_strand_frag = ( + ref.chrom, + ref.start, + ref.end, + ref.strand, + ) + + ref_strand = None + target_strand = None + if is_stranded: + ref_strand = ref_strand_frag + if output_table and output_table.strand_col: + target_strand = f'{table_name}."{output_table.strand_col}"' + + target_chrom_expr = f'{table_name}."{target_chrom}"' + target_start_expr = f'{table_name}."{target_start}"' + target_end_expr = f'{table_name}."{target_end}"' + + distance_expr = generate_distance_case( + ref_chrom, + ref_start, + ref_end, + ref_strand, + target_chrom_expr, + target_start_expr, + target_end_expr, + target_strand, + stranded=is_stranded, + signed=is_signed, + ) + abs_distance_expr = f"ABS({distance_expr})" + + where_clauses = [f"{ref_chrom} = {target_chrom_expr}"] + if is_stranded and ref_strand and target_strand: + where_clauses.append(f"{ref_strand} = {target_strand}") + if max_dist_value is not None: + where_clauses.append(f"({abs_distance_expr}) <= {max_dist_value}") + + return distance_expr, abs_distance_expr, where_clauses, passthrough + + +def _lateral_form( + expression: GIQLNearest, + ctx: ExpansionContext, + table_name: str, + target_ref: ResolvedRef, + ref: ResolvedInterval, +) -> exp.Expression: + """The portable LATERAL/standalone subquery. + + Builds a two-level subquery: an inner ``SELECT , AS + distance FROM WHERE ...`` that materializes the distance, wrapped by + an outer ``SELECT * FROM () AS x ORDER BY ABS(x.distance), x., + x. LIMIT k`` that orders on the *precomputed* ``distance`` column. For a + correlated placement the parent ``LATERAL`` correlates it to the outer row; + for a standalone (literal-reference) placement it stands alone. + + Splitting the distance computation (inner) from the ordering (outer) is + load-bearing for cross-engine support: + + * DuckDB's correlated-``LATERAL`` binder will not resolve a SELECT-list alias + named ``distance`` from inside an ``ORDER BY`` that also projects + ``.*``, so the order key must reference a *materialized* column + (``x.distance``) from the wrapping level rather than an alias in the same + SELECT. + * DataFusion's planner, given the distance ``CASE`` re-emitted inline in the + ``ORDER BY`` over the chromosome-equality prefiltered scan, rewrites the + filtered ``chrom`` to a self-comparison in one copy of the CASE but not the + other and trips ``SanityCheckPlan``; ordering on the materialized column + avoids re-deriving the key. + + A ``(start, end)`` tiebreaker follows ``ABS(distance)`` so rows tied at the + k-th distance order deterministically — and identically across engines and + against the decorrelated fallback's ranking — *when ``(start, end)`` + distinguishes the tied candidates*. Two target rows sharing both distance and + ``(start, end)`` remain order-ambiguous (no key here breaks that residual + tie); the two forms are set-equivalent up to such coordinate-duplicate ties + (#142 A5). + """ + k_value, *_ = _nearest_params(expression) + ( + distance_expr, + _abs_distance_expr, + where_clauses, + passthrough, + ) = _distance_and_filters(expression, table_name, target_ref, ref, ctx.capabilities) + where_sql = " AND ".join(where_clauses) + # The wrapping level reads the inner row's *bare* column names (the passthrough + # projected ``.*``), so the tiebreaker qualifies them by the wrapper + # alias, not the original ``table_name."col"``. + _chrom, target_start_col, target_end_col = target_ref.cols + wrapper = ctx.alias() + inner = ( + f"SELECT {passthrough}, {distance_expr} AS distance " + f"FROM {table_name} WHERE {where_sql}" + ) + sql = ( + f"(SELECT * FROM ({inner}) AS {wrapper} " + f'ORDER BY ABS({wrapper}."distance"), ' + f'{wrapper}."{target_start_col}", {wrapper}."{target_end_col}" ' + f"LIMIT {k_value})" + ) + return parse_one(sql, dialect=GIQLDialect) + + +def _outer_relation(ref: ResolvedInterval) -> tuple[str, str]: + """Return ``(physical_relation, alias)`` for the correlated reference table. + + The reference endpoints are alias-qualified fragments (``a."chrom"``). The + alias is the outer table's correlation name in the query; the physical + relation comes from the reference's backing :class:`~giql.table.Table`. Both + are needed to re-introduce the outer relation inside the decorrelated + subquery the fallback builds. + + The ``else alias`` branch (``ref.table is None``) only fires for a reference + whose backing table the resolver could not attach. For a *correlated* + NEAREST — the only shape that reaches the fallback — the reference is an + outer-row column, so the resolver always attaches its table and this branch + is not reached in practice. Falling back to ``alias`` (yielding a + cosmetically redundant ``FROM AS ``, where ``relation == + alias``) keeps the emitted SQL valid rather than emitting an empty relation + name should that invariant ever not hold; the assert pins the expectation. + """ + parsed = parse_one(ref.chrom, dialect=GIQLDialect) + alias = parsed.table if isinstance(parsed, exp.Column) else "" + if ref.table is not None: + relation = ref.table.name + else: + # Defensive fallback only: a correlated reference always carries a + # resolved backing table, so relation == alias here would be a redundant + # self-alias rather than a real two-name relation. + assert alias, ( + "correlated NEAREST fallback expected the reference to carry either a " + "backing table or an alias-qualified column" + ) + relation = alias + return relation, alias + + +def _projection_surfaces_reserved( + select: exp.Select, b_alias: str, reserved: list[str] +) -> bool: + """Whether *select*'s projection would surface the fallback's reserved columns. + + The decorrelated fallback exposes reserved rank/key columns on the join + relation *b_alias*. They reach the user output only through a projection that + stars over *b*: an unqualified ``*`` (``SELECT *``, which pulls every relation + including *b*) or a ``.*`` (matched case-insensitively unless the + qualifier is explicitly quoted, since engines fold unquoted identifiers). + Returns ``True`` for exactly those two shapes. + + The check must be exact. It must NOT match an explicit projection or a + ``.*`` (e.g. ``a.*``), because those never carry the reserved columns + and wrapping them in ``SELECT * EXCEPT ()`` is not caught at + transpile time — sqlglot builds it happily — but fails at engine runtime. An + unqualified ``*`` that already ``EXCEPT``s every reserved name is treated as + *not* surfacing, so a re-run (or a second finalizer) never double-wraps. + """ + for projection in select.expressions: + if isinstance(projection, exp.Star): + excepted = {col.name for col in (projection.args.get("except_") or [])} + if not set(reserved).issubset(excepted): + return True + elif isinstance(projection, exp.Column) and isinstance( + projection.this, exp.Star + ): + table = projection.args.get("table") + if table is None: + continue + # An unquoted ``B.*`` binds to the same relation as alias ``b`` because + # engines fold unquoted identifiers; match case-insensitively unless the + # qualifier is explicitly quoted (then its case is significant). + quoted = bool(table.args.get("quoted")) + if table.name == b_alias or ( + not quoted and table.name.casefold() == b_alias.casefold() + ): + return True + return False + + +def _wrap_star_except_reserved( + select: exp.Select, reserved: list[str], wrap_alias: str +) -> exp.Select: + """Wrap *select* in ``SELECT * EXCEPT (reserved...) FROM (