diff --git a/Cargo.lock b/Cargo.lock index 864aa8e5..040fcbdc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3728,6 +3728,7 @@ dependencies = [ "prost", "pyo3", "pyo3-build-config", + "sqlparser", "tokio", "tonic", ] diff --git a/Cargo.toml b/Cargo.toml index 482e10ee..d46a157d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,8 @@ prost = { version = "0.14", optional = true } # platform works on all CPython >= 3.10 (matching `requires-python`). Maturin # enables `pyo3/extension-module` through pyproject.toml for wheel builds; it # must stay disabled for ordinary Cargo test binaries so they link libpython. +# Parses SQL to share the CTEs and views a prepared query reads more than once. +sqlparser = { version = "0.62", features = ["visitor"] } pyo3 = { version = "0.29.0", features = ["abi3-py310"] } tokio = { version = "1.46.1", features = ["macros", "rt", "sync", "time"] } tonic = { version = "0.14", features = ["transport"], optional = true } diff --git a/docs/performance.md b/docs/performance.md index 60d7135a..a9c702a1 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -204,6 +204,15 @@ planning, as it is over tables of more than about a megabyte, a call plans the query again and runs as fast as `ctx.sql`. So does a call of a plan over a lazy table (Dask or a lazily opened file). +DataFusion plans a CTE or a view in full at each place a query reads +it. A model whose layers each read the layer before more than once +grows with every layer: a one-layer GPT's loss plans as 9,882 operators, +repeats of only 87 distinct subplans. `prepare` runs each CTE or view read more +than once just once per call, and keeps its rows in a table for the +rest of the query to read. That GPT's loss takes 16 ms per call, and +2.7 s with `ctx.sql`. A query with `WITH RECURSIVE`, or a subquery with +CTEs of its own, is planned as it's written. + ## Round-trip faster with ORDER BY Results that arrive **grid-ordered** — sorted by the dimension columns, diff --git a/src/lib.rs b/src/lib.rs index 003d8707..fa077033 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,7 @@ #[cfg(feature = "flight")] mod flight; +mod share; use std::collections::{HashMap, HashSet}; use std::ffi::CString; @@ -1656,6 +1657,7 @@ impl LazyArrowStreamTable { #[pymodule] fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_function(wrap_pyfunction!(share::share, m)?)?; #[cfg(feature = "flight")] m.add_class::()?; Ok(()) diff --git a/src/share.rs b/src/share.rs new file mode 100644 index 00000000..0d37d8c3 --- /dev/null +++ b/src/share.rs @@ -0,0 +1,301 @@ +//! Splitting a query so each CTE or view it reads more than once runs once. +//! +//! DataFusion plans a CTE or a view in full at every place a query reads it, +//! so a model whose layers each read the one before several times plans and +//! runs the first layers many times over. [`share`] rewrites the query into +//! statements: one for each CTE or view read more than once, which the caller +//! stores as a table under the name the statement gives, and then the query +//! itself, reading those tables. + +use std::collections::{HashMap, HashSet}; +use std::ops::ControlFlow; + +use pyo3::prelude::*; +use sqlparser::ast::helpers::attached_token::AttachedToken; +use sqlparser::ast::{ + Cte, Ident, ObjectName, ObjectNamePart, Query, Statement, TableAlias, TableAliasColumnDef, + TableFactor, VisitMut, VisitorMut, With, +}; +use sqlparser::dialect::GenericDialect; +use sqlparser::parser::Parser; + +/// A CTE or view, hoisted out of the query that defined it. +struct Def { + query: Query, + columns: Vec, + /// The definitions it reads, once for each place it reads them. + reads: Vec, +} + +/// Why a query is left as it is. +struct Unsupported; + +struct Hoister<'a> { + prefix: &'a str, + views: &'a HashMap, + defs: Vec, + view_defs: HashMap, + expanding: HashSet, +} + +/// The name a query's identifier resolves to: unquoted identifiers are +/// lowercased, as DataFusion does. +fn normalized(ident: &Ident) -> String { + match ident.quote_style { + Some(_) => ident.value.clone(), + None => ident.value.to_lowercase(), + } +} + +impl Hoister<'_> { + fn name(&self, index: usize) -> String { + format!("{}{index}", self.prefix) + } + + /// Hoists `query`'s own CTEs into definitions, and renames each table + /// it reads that is a CTE in `scope` or a view to that definition's name. + /// Returns the definitions it reads. + fn hoist( + &mut self, + query: &mut Query, + scope: &HashMap, + ) -> Result, Unsupported> { + let mut scope = scope.clone(); + if let Some(with) = query.with.take() { + if with.recursive { + return Err(Unsupported); + } + for cte in with.cte_tables { + if cte.materialized.is_some() || cte.from.is_some() { + return Err(Unsupported); + } + let mut body = *cte.query; + let reads = self.hoist(&mut body, &scope)?; + self.defs.push(Def { + query: body, + columns: cte.alias.columns, + reads, + }); + scope.insert(normalized(&cte.alias.name), self.defs.len() - 1); + } + } + let mut renamer = Renamer { + hoister: self, + scope: &scope, + top: true, + reads: Vec::new(), + }; + match query.visit(&mut renamer) { + ControlFlow::Break(e) => Err(e), + ControlFlow::Continue(()) => Ok(renamer.reads), + } + } + + /// The definition of the view `name`, hoisting it the first time. + fn view(&mut self, name: &str) -> Result { + if let Some(&index) = self.view_defs.get(name) { + return Ok(index); + } + if !self.expanding.insert(name.to_string()) { + return Err(Unsupported); + } + let sql = &self.views[name]; + let mut statements = Parser::parse_sql(&GenericDialect {}, sql).map_err(|_| Unsupported)?; + let Some(Statement::CreateView(view)) = statements.pop() else { + return Err(Unsupported); + }; + if !statements.is_empty() || !view.columns.is_empty() { + return Err(Unsupported); + } + let mut body = *view.query; + // A view reads the catalog, not the CTEs of the query that reads it. + let reads = self.hoist(&mut body, &HashMap::new())?; + self.defs.push(Def { + query: body, + columns: Vec::new(), + reads, + }); + self.expanding.remove(name); + let index = self.defs.len() - 1; + self.view_defs.insert(name.to_string(), index); + Ok(index) + } +} + +/// Renames the tables a query reads that are CTEs in scope, or views. +struct Renamer<'a, 'b> { + hoister: &'a mut Hoister<'b>, + scope: &'a HashMap, + /// Whether the next query visited is the one being renamed, whose CTEs + /// were hoisted already, rather than a subquery of it. + top: bool, + reads: Vec, +} + +impl VisitorMut for Renamer<'_, '_> { + type Break = Unsupported; + + fn pre_visit_query(&mut self, query: &mut Query) -> ControlFlow { + // A subquery's own CTEs would shadow names in scope; leave such + // queries to DataFusion. + if !std::mem::take(&mut self.top) && query.with.is_some() { + return ControlFlow::Break(Unsupported); + } + ControlFlow::Continue(()) + } + + fn pre_visit_table_factor(&mut self, factor: &mut TableFactor) -> ControlFlow { + let TableFactor::Table { + name, alias, args, .. + } = factor + else { + return ControlFlow::Continue(()); + }; + let [ObjectNamePart::Identifier(ident)] = name.0.as_slice() else { + return ControlFlow::Continue(()); + }; + if args.is_some() { + return ControlFlow::Continue(()); + } + let key = normalized(ident); + let index = match self.scope.get(&key) { + Some(&index) => index, + None if self.hoister.views.contains_key(&key) => match self.hoister.view(&key) { + Ok(index) => index, + Err(e) => return ControlFlow::Break(e), + }, + None => return ControlFlow::Continue(()), + }; + // Columns qualified by the name read, as in `x0.val`, still resolve. + if alias.is_none() { + *alias = Some(TableAlias { + explicit: true, + name: ident.clone(), + columns: Vec::new(), + at: None, + }); + } + *name = ObjectName::from(vec![Ident::new(self.hoister.name(index))]); + self.reads.push(index); + ControlFlow::Continue(()) + } +} + +/// `query` with a `WITH` of the definitions it reads that aren't `shared`, +/// read through those that aren't either. +fn with_inlined( + mut query: Query, + reads: &[usize], + defs: &[Def], + shared: &[bool], + hoister: &Hoister, +) -> Query { + let mut needed = HashSet::new(); + let mut stack: Vec = reads.iter().copied().filter(|&i| !shared[i]).collect(); + while let Some(i) = stack.pop() { + if needed.insert(i) { + stack.extend(defs[i].reads.iter().copied().filter(|&j| !shared[j])); + } + } + let mut needed: Vec = needed.into_iter().collect(); + // A definition only reads those made before it. + needed.sort_unstable(); + if !needed.is_empty() { + query.with = Some(With { + with_token: AttachedToken::empty(), + recursive: false, + cte_tables: needed + .into_iter() + .map(|i| Cte { + alias: TableAlias { + explicit: false, + name: Ident::new(hoister.name(i)), + columns: defs[i].columns.clone(), + at: None, + }, + query: Box::new(defs[i].query.clone()), + from: None, + materialized: None, + closing_paren_token: AttachedToken::empty(), + }) + .collect(), + }); + } + query +} + +/// Splits `query` so that each CTE or view it reads more than once runs once. +/// +/// `views` maps the name of each view in the default schema to the SQL that +/// created it. Returns `(steps, query)`: each step is a table name and the +/// SQL whose rows that table must hold before the next step runs, and +/// `query` reads them. Names start with `prefix`. Returns `None` when +/// nothing is read twice, or for SQL this doesn't rewrite: a recursive CTE, +/// a subquery with CTEs of its own, or more than one statement. +#[pyfunction] +pub fn share( + query: &str, + views: HashMap, + prefix: &str, +) -> Option<(Vec<(String, String)>, String)> { + let mut statements = Parser::parse_sql(&GenericDialect {}, query).ok()?; + let Some(Statement::Query(mut query)) = statements.pop() else { + return None; + }; + if !statements.is_empty() { + return None; + } + let mut hoister = Hoister { + prefix, + views: &views, + defs: Vec::new(), + view_defs: HashMap::new(), + expanding: HashSet::new(), + }; + let reads = hoister.hoist(&mut query, &HashMap::new()).ok()?; + + // How often each definition runs if every other is inlined: once for + // each place it's read, counting only places in definitions that run. + let defs = std::mem::take(&mut hoister.defs); + let mut used = vec![false; defs.len()]; + let mut stack = reads.clone(); + while let Some(i) = stack.pop() { + if !std::mem::replace(&mut used[i], true) { + stack.extend(&defs[i].reads); + } + } + let mut times = vec![0usize; defs.len()]; + for i in reads.iter().copied().chain( + (0..defs.len()) + .filter(|&i| used[i]) + .flat_map(|i| defs[i].reads.iter().copied()), + ) { + times[i] += 1; + } + // Stored once, a definition read twice runs once; every other is then + // inlined into the one place it's read, which runs once too. + let shared: Vec = times.iter().map(|&n| n > 1).collect(); + if !shared.contains(&true) { + return None; + } + + let steps = (0..defs.len()) + .filter(|&i| shared[i]) + .map(|i| { + let name = hoister.name(i); + let mut not_shared = shared.clone(); + not_shared[i] = false; + let select = Parser::parse_sql(&GenericDialect {}, &format!("SELECT * FROM {name}")) + .ok() + .and_then(|mut s| match s.pop() { + Some(Statement::Query(q)) => Some(*q), + _ => None, + }) + .expect("a SELECT of one table parses"); + let step = with_inlined(select, &[i], &defs, ¬_shared, &hoister); + (name, step.to_string()) + }) + .collect(); + let query = with_inlined(*query, &reads, &defs, &shared, &hoister); + Some((steps, query.to_string())) +} diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 8cc356bc..51c2c235 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -262,3 +262,108 @@ def test_a_quoted_table_replaced_in_sql_counts_too(): ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect() with pytest.raises(ValueError, match="registered again"): query() + + +def _planned(query, a, b, *setup): + """The rows of ``query`` planned from scratch over ``a`` and ``b``.""" + fresh = Context() + fresh.from_dataset("a", a) + fresh.from_dataset("b", b) + for statement in setup: + fresh.sql(statement) + return fresh.sql(query).to_arrow_table() + + +def _calls(query, *setup): + """A prepared ``query``'s rows on new values of ``a``, against planning + it.""" + rng = np.random.default_rng(1) + b = _b(rng.normal(size=(4, 3))) + ctx = Context() + ctx.from_dataset("a", _a(rng.normal(size=(4, 3)))) + ctx.from_dataset("b", b) + for statement in setup: + ctx.sql(statement) + prepared = ctx.prepare(query) + for _ in range(3): + a = _a(rng.normal(size=(4, 3))) + yield prepared(a=a), _planned(query, a, b, *setup) + + +# Each reads `t` twice, so a call runs `t` once and stores its rows. +@pytest.mark.parametrize( + "query", + [ + # A CTE read twice. + ( + "WITH t AS (SELECT i, j, 2 * v AS v2 FROM a) " + "SELECT x.i, x.j, x.v2 * y.v2 AS p FROM t x JOIN t y " + "ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j" + ), + # Read twice, by a CTE read twice itself. + ( + "WITH t AS (SELECT i, j, v FROM a), " + "u AS (SELECT x.i, x.j, x.v + y.v AS v FROM t x JOIN t y " + "ON x.i = y.i AND x.j = y.j), " + "s AS (SELECT u.i, SUM(u.v * b.w) AS s FROM u JOIN b " + "ON u.i = b.i AND u.j = b.j GROUP BY u.i) " + "SELECT p.i, p.s + q.s AS s FROM s p JOIN s q ON p.i = q.i ORDER BY p.i" + ), + # Rows whose number changes from call to call. + ( + "WITH t AS (SELECT i, j, v FROM a WHERE v > 0) " + "SELECT (SELECT COUNT(*) FROM t) AS n, (SELECT SUM(v) FROM t) AS s" + ), + # Rows with nulls, and with strings, which aren't written in place. + ( + "WITH t AS (SELECT i, j, CASE WHEN v > 0 THEN v END AS v, " + "CASE WHEN v > 0 THEN 'up' ELSE 'down' END AS d FROM a) " + "SELECT x.i, x.j, x.v, y.d FROM t x JOIN t y " + "ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j" + ), + # Columns named where the CTE is. + ( + "WITH t(r, c, val) AS (SELECT i, j, v FROM a) " + "SELECT x.r, x.c, x.val + y.val AS z FROM t x JOIN t y " + "ON x.r = y.r AND x.c = y.c ORDER BY x.r, x.c" + ), + ], +) +def test_a_cte_read_more_than_once_gives_the_rows_planning_gives(query): + for rows, want in _calls(query): + assert rows.equals(want) + + +def test_a_view_read_more_than_once_gives_the_rows_planning_gives(): + view = ( + "CREATE VIEW t AS WITH p AS (SELECT i, j, 2 * v AS v FROM a) " + "SELECT x.i, x.j, x.v + y.v AS v FROM p x JOIN p y " + "ON x.i = y.i AND x.j = y.j" + ) + query = ( + "SELECT x.i, SUM(x.v * y.v) AS s FROM t x JOIN t y " + "ON x.i = y.i AND x.j = y.j GROUP BY x.i ORDER BY x.i" + ) + for rows, want in _calls(query, view): + assert rows.equals(want) + + +def test_a_cte_named_like_a_view_hides_it(): + view = "CREATE VIEW t AS SELECT i, j, -v AS v FROM a" + query = ( + "WITH t AS (SELECT i, j, v FROM a) " + "SELECT x.i, x.j, x.v + y.v AS v FROM t x JOIN t y " + "ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j" + ) + for rows, want in _calls(query, view): + assert rows.equals(want) + + +def test_a_recursive_cte_runs_again(): + query = ( + "WITH RECURSIVE r(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM r " + "WHERE n < 4) SELECT SUM(n) AS s FROM r" + ) + prepared = Context().prepare(query) + assert prepared().to_pydict() == {"s": [10]} + assert prepared().to_pydict() == {"s": [10]} diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 0c56d039..2ed3b0d7 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -12,6 +12,7 @@ from types import ModuleType from . import cftime as cft +from ._native import share as _share_sql from .df import ( Chunks, TableNames, @@ -380,45 +381,35 @@ class PreparedQuery: To change anything else, register the new Dataset and prepare the query again. A call raises if a table the query reads has been registered again since it was prepared. Functions such as ``now()`` - are evaluated once, when the query is prepared. + are evaluated once, when the query is prepared. A CTE or view the + query reads more than once runs once per call. """ def __init__(self, ctx: Context, query: str): self._ctx = ctx - start = time.perf_counter() - df = SessionContext.sql(ctx, query) - plan = df.execution_plan() - planning = time.perf_counter() - start - self._schema = df.schema() - # Planned again from its optimized form, the query reads the tables - # it was prepared with, and functions such as now() keep the values - # they were folded into. - self._optimized = df.optimized_logical_plan() + # Each CTE or view read more than once runs as a step of its own, + # whose rows are stored, under a name of this query's, as a table + # the later steps and the query read. + shared = _share(ctx, query, f"_xql_{uuid.uuid4().hex}_") + steps, query = shared if shared is not None else ([], query) + self._steps = [] + for name, sql in steps: + step = _Statement(ctx, sql) + self._steps.append((step, _Stored(ctx, name, step()))) + self._query = _Statement(ctx, query) + stored = {name for name, _ in steps} # Each table the query reads, by how often its name had been # deregistered; and of those, the Datasets a call may write into. self._read = { name: ctx._deregistrations.get(name, 0) - for name in _tables_read(df.logical_plan()) + for statement in [step for step, _ in self._steps] + [self._query] + for name in statement.reads - stored } self._registered = { name: ctx._registered_datasets[name] for name in self._read if name in ctx._registered_datasets } - # A plan keeps state from a run (a hash join keeps the table it - # built), so each call runs a fresh copy of it, decoded from its - # serialized form. Copying a plan copies its in-memory tables, so - # over large ones it's slower than planning again; a plan that reads - # a lazy table, or is nested deeper than protobuf decodes, can't be - # copied at all. Either way, each call plans the query again. - self._plan: ExecutionPlan | None = None - start = time.perf_counter() - try: - _copy(ctx, plan) - except Exception: # noqa: BLE001 - DataFusion raises a bare Exception - return - if time.perf_counter() - start < planning: - self._plan = plan def __call__(self, **tables: xr.Dataset) -> pa.Table: """Run the query, after writing each Dataset given into the table @@ -443,6 +434,55 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table: "prepared; prepare the query again" ) self._ctx._write(tables) + for i, (step, stored) in enumerate(self._steps): + rows = step() + if not stored.write(rows): + # The rows don't fit the table the later statements were + # planned over, so store them in a new one and plan those + # statements again. + stored.replace(rows) + for j in range(i + 1, len(self._steps)): + later, stored = self._steps[j] + later = _Statement(self._ctx, later.sql) + stored.replace(later()) + self._steps[j] = (later, stored) + self._query = _Statement(self._ctx, self._query.sql) + break + return self._query() + + +class _Statement: + """One query, planned once, to run many times.""" + + def __init__(self, ctx: SessionContext, sql: str): + self._ctx = ctx + self.sql = sql + start = time.perf_counter() + df = SessionContext.sql(ctx, sql) + plan = df.execution_plan() + planning = time.perf_counter() - start + self._schema = df.schema() + # Planned again from its optimized form, the query reads the tables + # it was prepared with, and functions such as now() keep the values + # they were folded into. + self._optimized = df.optimized_logical_plan() + self.reads = _tables_read(df.logical_plan()) + # A plan keeps state from a run (a hash join keeps the table it + # built), so each call runs a fresh copy of it, decoded from its + # serialized form. Copying a plan copies its in-memory tables, so + # over large ones it's slower than planning again; a plan that reads + # a lazy table, or is nested deeper than protobuf decodes, can't be + # copied at all. Either way, each call plans the query again. + self._plan: ExecutionPlan | None = None + start = time.perf_counter() + try: + _copy(ctx, plan) + except Exception: # noqa: BLE001 - DataFusion raises a bare Exception + return + if time.perf_counter() - start < planning: + self._plan = plan + + def __call__(self) -> pa.Table: if self._plan is None: df = self._ctx.create_dataframe_from_logical_plan(self._optimized) return df.to_arrow_table() @@ -455,6 +495,82 @@ def __call__(self, **tables: xr.Dataset) -> pa.Table: return pa.Table.from_batches(batches, self._schema) +class _Stored: + """A step's rows, stored as a table that the statements after it read. + + Its numeric columns share memory with arrays it owns, so a step's new + rows can be written into the table in place. + """ + + def __init__(self, ctx: SessionContext, name: str, rows: pa.Table): + self._ctx = ctx + self._name = name + self.replace(rows) + + def replace(self, rows: pa.Table) -> None: + """Store ``rows`` as a new table under the name.""" + self._schema = rows.schema + self._rows = rows.num_rows + self._arrays = [_writable(column) for column in rows.columns] + columns = [ + pa.array(array) if array is not None else column.combine_chunks() + for array, column in zip(self._arrays, rows.columns) + ] + batch = pa.RecordBatch.from_arrays(columns, schema=self._schema) + SessionContext.deregister_table(self._ctx, self._name) + SessionContext.register_record_batches(self._ctx, self._name, [[batch]]) + + def write(self, rows: pa.Table) -> bool: + """Write ``rows`` into the table in place, if they fit it: the same + schema and number of rows, in numeric columns without nulls.""" + if rows.schema != self._schema or rows.num_rows != self._rows: + return False + arrays = [array for array in self._arrays if array is not None] + if len(arrays) < len(self._arrays) or any( + column.null_count for column in rows.columns + ): + return False + for array, column in zip(arrays, rows.columns): + array[...] = column.to_numpy() + return True + + +def _writable(column: pa.ChunkedArray) -> np.ndarray | None: + """A copy of ``column`` that pyarrow can share memory with, if it's + numeric and has no nulls.""" + numeric = pa.types.is_integer(column.type) or pa.types.is_floating( + column.type + ) + if not numeric or column.null_count: + return None + return np.array(column.to_numpy(), copy=True) + + +def _share( + ctx: SessionContext, query: str, prefix: str +) -> tuple[list[tuple[str, str]], str] | None: + """``query`` as steps that each run a CTE or view it reads more than + once, and the query reading what they store: see ``_native.share``.""" + try: + views = SessionContext.sql( + ctx, + "SELECT v.table_name, v.definition FROM information_schema.views v" + " JOIN information_schema.df_settings c" + " ON c.name = 'datafusion.catalog.default_catalog'" + " AND v.table_catalog = c.value" + " JOIN information_schema.df_settings s" + " ON s.name = 'datafusion.catalog.default_schema'" + " AND v.table_schema = s.value" + " WHERE v.definition IS NOT NULL", + ).to_pydict() + except Exception: # noqa: BLE001 - no information schema to read + views = {"table_name": [], "definition": []} + shared: tuple[list[tuple[str, str]], str] | None = _share_sql( + query, dict(zip(views["table_name"], views["definition"])), prefix + ) + return shared + + def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan: """A copy of ``plan`` with none of its state from a run, reading its in-memory tables' values as they are now."""