@@ -262,3 +262,108 @@ def test_a_quoted_table_replaced_in_sql_counts_too():
262262 ctx .sql ('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)' ).collect ()
263263 with pytest .raises (ValueError , match = "registered again" ):
264264 query ()
265+
266+
267+ def _planned (query , a , b , * setup ):
268+ """The rows of ``query`` planned from scratch over ``a`` and ``b``."""
269+ fresh = Context ()
270+ fresh .from_dataset ("a" , a )
271+ fresh .from_dataset ("b" , b )
272+ for statement in setup :
273+ fresh .sql (statement )
274+ return fresh .sql (query ).to_arrow_table ()
275+
276+
277+ def _calls (query , * setup ):
278+ """A prepared ``query``'s rows on new values of ``a``, against planning
279+ it."""
280+ rng = np .random .default_rng (1 )
281+ b = _b (rng .normal (size = (4 , 3 )))
282+ ctx = Context ()
283+ ctx .from_dataset ("a" , _a (rng .normal (size = (4 , 3 ))))
284+ ctx .from_dataset ("b" , b )
285+ for statement in setup :
286+ ctx .sql (statement )
287+ prepared = ctx .prepare (query )
288+ for _ in range (3 ):
289+ a = _a (rng .normal (size = (4 , 3 )))
290+ yield prepared (a = a ), _planned (query , a , b , * setup )
291+
292+
293+ # Each reads `t` twice, so a call runs `t` once and stores its rows.
294+ @pytest .mark .parametrize (
295+ "query" ,
296+ [
297+ # A CTE read twice.
298+ (
299+ "WITH t AS (SELECT i, j, 2 * v AS v2 FROM a) "
300+ "SELECT x.i, x.j, x.v2 * y.v2 AS p FROM t x JOIN t y "
301+ "ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j"
302+ ),
303+ # Read twice, by a CTE read twice itself.
304+ (
305+ "WITH t AS (SELECT i, j, v FROM a), "
306+ "u AS (SELECT x.i, x.j, x.v + y.v AS v FROM t x JOIN t y "
307+ "ON x.i = y.i AND x.j = y.j), "
308+ "s AS (SELECT u.i, SUM(u.v * b.w) AS s FROM u JOIN b "
309+ "ON u.i = b.i AND u.j = b.j GROUP BY u.i) "
310+ "SELECT p.i, p.s + q.s AS s FROM s p JOIN s q ON p.i = q.i ORDER BY p.i"
311+ ),
312+ # Rows whose number changes from call to call.
313+ (
314+ "WITH t AS (SELECT i, j, v FROM a WHERE v > 0) "
315+ "SELECT (SELECT COUNT(*) FROM t) AS n, (SELECT SUM(v) FROM t) AS s"
316+ ),
317+ # Rows with nulls, and with strings, which aren't written in place.
318+ (
319+ "WITH t AS (SELECT i, j, CASE WHEN v > 0 THEN v END AS v, "
320+ "CASE WHEN v > 0 THEN 'up' ELSE 'down' END AS d FROM a) "
321+ "SELECT x.i, x.j, x.v, y.d FROM t x JOIN t y "
322+ "ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j"
323+ ),
324+ # Columns named where the CTE is.
325+ (
326+ "WITH t(r, c, val) AS (SELECT i, j, v FROM a) "
327+ "SELECT x.r, x.c, x.val + y.val AS z FROM t x JOIN t y "
328+ "ON x.r = y.r AND x.c = y.c ORDER BY x.r, x.c"
329+ ),
330+ ],
331+ )
332+ def test_a_cte_read_more_than_once_gives_the_rows_planning_gives (query ):
333+ for rows , want in _calls (query ):
334+ assert rows .equals (want )
335+
336+
337+ def test_a_view_read_more_than_once_gives_the_rows_planning_gives ():
338+ view = (
339+ "CREATE VIEW t AS WITH p AS (SELECT i, j, 2 * v AS v FROM a) "
340+ "SELECT x.i, x.j, x.v + y.v AS v FROM p x JOIN p y "
341+ "ON x.i = y.i AND x.j = y.j"
342+ )
343+ query = (
344+ "SELECT x.i, SUM(x.v * y.v) AS s FROM t x JOIN t y "
345+ "ON x.i = y.i AND x.j = y.j GROUP BY x.i ORDER BY x.i"
346+ )
347+ for rows , want in _calls (query , view ):
348+ assert rows .equals (want )
349+
350+
351+ def test_a_cte_named_like_a_view_hides_it ():
352+ view = "CREATE VIEW t AS SELECT i, j, -v AS v FROM a"
353+ query = (
354+ "WITH t AS (SELECT i, j, v FROM a) "
355+ "SELECT x.i, x.j, x.v + y.v AS v FROM t x JOIN t y "
356+ "ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j"
357+ )
358+ for rows , want in _calls (query , view ):
359+ assert rows .equals (want )
360+
361+
362+ def test_a_recursive_cte_runs_again ():
363+ query = (
364+ "WITH RECURSIVE r(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM r "
365+ "WHERE n < 4) SELECT SUM(n) AS s FROM r"
366+ )
367+ prepared = Context ().prepare (query )
368+ assert prepared ().to_pydict () == {"s" : [10 ]}
369+ assert prepared ().to_pydict () == {"s" : [10 ]}
0 commit comments