From 938d2ed53f7c638a39998e05280852c116276875 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 26 Oct 2026 10:58:45 -0700 Subject: [PATCH 1/6] Let the dense algorithm check that a domain's keys are unique A plan that joins a fold back to its domain, then groups again by the domain's keys, is the fold from its domain exactly when no two domain rows reach one group. Options::unique_domain makes fold_dense decline otherwise. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-kernels/src/dense.rs | 52 ++++++++---- crates/einfold-kernels/src/lib.rs | 13 ++- .../einfold-kernels/tests/kernels_vs_sql.rs | 82 +++++++++++++++++++ 3 files changed, 128 insertions(+), 19 deletions(-) diff --git a/crates/einfold-kernels/src/dense.rs b/crates/einfold-kernels/src/dense.rs index 841cf0b3..0d064fc5 100644 --- a/crates/einfold-kernels/src/dense.rs +++ b/crates/einfold-kernels/src/dense.rs @@ -136,7 +136,8 @@ fn par_tasks(tasks: usize, threads: usize, f: impl Fn(usize) -> T + Syn /// The fold of `a` and `b` (and the batches of `domain`, if the fold has one) /// by the dense algorithm on `threads` threads, or `None` when its conditions /// don't hold and the hash path must be used instead. With a domain, there is -/// one output batch per domain batch. +/// one output batch per domain batch. With `unique_domain`, it is also `None` +/// when two of the domain's rows locate the same cell of the output. pub(crate) fn fold( fold: &Fold, a: &RecordBatch, @@ -144,6 +145,7 @@ pub(crate) fn fold( domain: Option<&[RecordBatch]>, schema: SchemaRef, threads: usize, + unique_domain: bool, ) -> Result>> { let [op_a, op_b] = fold.operands() else { return Ok(None); @@ -292,10 +294,18 @@ pub(crate) fn fold( strides: &strides, extent: &extent, }; - let batches = domain - .iter() - .map(|f| domain_batch(f, spec, &cells, &value, &schema)) - .collect::>()?; + let mut seen = vec![false; if unique_domain { nx * ny } else { 0 }]; + let mut batches = Vec::with_capacity(domain.len()); + for f in domain { + let cell = locate(f, &cells)?; + if unique_domain + && (cell.iter()) + .any(|&c| c != NO_GROUP && std::mem::replace(&mut seen[c], true)) + { + return Ok(None); + } + batches.push(domain_batch(f, &cell, spec, &value, &schema)?); + } Ok(Some(batches)) } (None, None) => { @@ -346,18 +356,13 @@ struct Cells<'a> { extent: &'a [usize], } -/// The output rows for one batch of the domain: its carried columns, and the -/// value of the cell its keys locate. A NULL key, or one outside every -/// operand's range, matches no group (the operands have no NULL keys), so its -/// row gets the default, under either key equality. -fn domain_batch( - batch: &RecordBatch, - spec: &Domain, - cells: &Cells, - value: &impl Fn(usize) -> Option, - schema: &SchemaRef, -) -> Result { - const NO_GROUP: usize = usize::MAX; +/// A domain row whose keys locate no cell of the output. +const NO_GROUP: usize = usize::MAX; + +/// The cell of the output that each row of a domain batch's keys locate, or +/// [`NO_GROUP`]. A NULL key, or one outside every operand's range, matches no +/// group (the operands have no NULL keys), under either key equality. +fn locate(batch: &RecordBatch, cells: &Cells) -> Result> { let mut cell = vec![0usize; batch.num_rows()]; for ((&d, &stride), &ColumnIndex(col)) in cells.dims.iter().zip(cells.strides).zip(cells.keys) { // Out-of-range unsigned keys become NULL here, so match no group. @@ -375,6 +380,19 @@ fn domain_batch( }; } } + Ok(cell) +} + +/// The output rows for one batch of the domain: its carried columns, and the +/// value of the cell each row locates (see [`locate`]), or the default where +/// it locates none. +fn domain_batch( + batch: &RecordBatch, + cell: &[usize], + spec: &Domain, + value: &impl Fn(usize) -> Option, + schema: &SchemaRef, +) -> Result { let at = |c: usize| { if c == NO_GROUP { spec.missing diff --git a/crates/einfold-kernels/src/lib.rs b/crates/einfold-kernels/src/lib.rs index 56bf16b8..2aadd726 100644 --- a/crates/einfold-kernels/src/lib.rs +++ b/crates/einfold-kernels/src/lib.rs @@ -61,6 +61,11 @@ pub struct Options { pub host: Host, /// The most threads to use. The result's bits don't depend on it. pub threads: usize, + /// Decline unless no two of the domain's rows locate the same group. A + /// plan that joins a fold back to its domain and groups again by the + /// domain's keys is the fold from its domain exactly when they are + /// unique. Only [`fold_dense`] checks it, as it can decline. + pub unique_domain: bool, } /// Which algorithm computed a fold. @@ -195,7 +200,8 @@ pub fn fold_dense( schema: SchemaRef, options: Options, ) -> Result>> { - dense::fold(fold, a, b, domain, schema, options.threads.max(1)) + let threads = options.threads.max(1); + dense::fold(fold, a, b, domain, schema, threads, options.unique_domain) } /// `fold` over operands `a` and `b`, and the batches of its domain if it has @@ -213,8 +219,11 @@ pub fn fold( schema: SchemaRef, options: Options, ) -> Result<(Algorithm, Vec)> { + if options.unique_domain { + return bad("only fold_dense checks that a domain's keys are unique".into()); + } let threads = options.threads.max(1); - if let Some(batches) = dense::fold(fold, a, b, domain, schema.clone(), threads)? { + if let Some(batches) = dense::fold(fold, a, b, domain, schema.clone(), threads, false)? { return Ok((Algorithm::Dense, batches)); } let batches = hash::fold(fold, a, b, domain, schema, options.host)?; diff --git a/crates/einfold-kernels/tests/kernels_vs_sql.rs b/crates/einfold-kernels/tests/kernels_vs_sql.rs index e4bec40f..67a965ee 100644 --- a/crates/einfold-kernels/tests/kernels_vs_sql.rs +++ b/crates/einfold-kernels/tests/kernels_vs_sql.rs @@ -23,6 +23,10 @@ //! that of the exact sum, and so is the host's. //! 3. **Deterministic.** The same rows, in another order and other batches, //! on another number of threads, give the same rows, bit for bit. +//! 4. **Unique domains.** Where the dense path runs with a domain, asking it +//! to check that the domain's keys are unique declines whenever two domain +//! rows have the keys of one group, and otherwise, when it accepts, gives +//! the same rows, bit for bit. //! //! Half the cases are drawn from a *dense family*, shaped so that the dense //! algorithm may apply: non-negative integer keys with no NULLs in the @@ -635,12 +639,84 @@ fn run_fold(case: &Case, rng: &mut Rng) -> (Vec, Algorithm) { let options = Options { host: HOST, threads: 1 + rng.below(8), + unique_domain: false, }; let (algorithm, out) = einfold_kernels::fold(&case.fold, &a, &b, domain.as_deref(), schema, options).unwrap(); (out, algorithm) } +/// Checks invariant 4 for a case with a domain that the dense path computes, +/// and says whether the checked path accepted. +fn check_unique_domain(case: &Case, rng: &mut Rng, fail: impl Fn(String)) -> bool { + let fold = &case.fold; + let domain = case.domain.as_ref().expect("a case with a domain"); + let whole = |t: &Table, rng: &mut Rng| { + let batches: Vec = t.partitions(rng).into_iter().flatten().collect(); + concat_batches(&t.schema, &batches).unwrap() + }; + let (a, b) = (whole(&case.a, rng), whole(&case.b, rng)); + let d: Vec = domain.partitions(rng).into_iter().flatten().collect(); + let operands = [case.a.schema.as_ref(), case.b.schema.as_ref()]; + let mut dense = |fold: &Fold, d: Option<&[RecordBatch]>, unique_domain| { + let schema = einfold_kernels::schema(fold, &operands, d.map(|_| domain.schema.as_ref())); + let options = Options { + host: HOST, + threads: 1 + rng.below(8), + unique_domain, + }; + einfold_kernels::fold_dense(fold, &a, &b, d, schema.unwrap(), options).unwrap() + }; + let keys = fold.outputs().len(); + let unchecked = dense(fold, Some(&d), false).expect("the dense path applies"); + // The groups, from the same fold without its domain. + let outputs = fold + .outputs() + .iter() + .map(|o| Output { + domain_key: None, + ..*o + }) + .collect(); + let bare = Fold::try_new( + fold.dimensions().to_vec(), + fold.operands().to_vec(), + outputs, + fold.row(), + fold.aggregate().clone(), + None, + ) + .unwrap(); + let groups: std::collections::HashSet = match dense(&bare, None, false) { + Some(out) => rows(&out, keys).into_iter().map(|(k, _)| k).collect(), + None => return false, + }; + let mut seen = std::collections::HashSet::new(); + let repeated = rows(&unchecked, keys) + .into_iter() + .any(|(k, _)| groups.contains(&k) && !seen.insert(k)); + match dense(fold, Some(&d), true) { + Some(_) if repeated => { + fail("the checked dense path accepted two domain rows with one group's keys".into()); + true + } + Some(checked) => { + let bits = |out: &[RecordBatch]| { + let mut r: Vec<_> = (rows(out, keys).into_iter()) + .map(|(k, v)| (k, v.map(f64::to_bits))) + .collect(); + r.sort(); + r + }; + if bits(&checked) != bits(&unchecked) { + fail("the checked dense path changed the output".into()); + } + true + } + None => false, + } +} + #[tokio::test] async fn fold_matches_sql() { let seed: u64 = std::env::var("EINFOLD_SEED") @@ -653,6 +729,7 @@ async fn fold_matches_sql() { Err(_) => 400, }; let (mut dense_family, mut dense_ran) = (0, 0); + let (mut unique_checked, mut unique_accepted) = (0, 0); for i in 0..cases { let seed = seed.wrapping_add(i); let mut rng = Rng::new(seed); @@ -667,6 +744,10 @@ async fn fold_matches_sql() { let (again, _) = run_fold(&case, &mut rng); dense_family += usize::from(dense); dense_ran += usize::from(ran == Algorithm::Dense); + if ran == Algorithm::Dense && case.domain.is_some() { + unique_checked += 1; + unique_accepted += usize::from(check_unique_domain(&case, &mut rng, |w| fail(w))); + } let (ours, again) = (rows(&ours, keys), rows(&again, keys)); // SQL doesn't order rows, so the outputs are compared as multisets. let bits = |r: &[Row]| -> Vec<(String, Option)> { @@ -705,6 +786,7 @@ async fn fold_matches_sql() { } } eprintln!("dense path: {dense_ran} of {cases} cases; dense family: {dense_family}"); + eprintln!("unique-domain check: accepted {unique_accepted} of {unique_checked}"); // About a third of the dense family qualifies (the rest has duplicate // tuples, keys too large for BIGINT, or holes where existence matters); // if far fewer did, the dense path would go untested. From 31b2e6a570d3e68d53ff2c62f35568cecf267271 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 26 Oct 2026 11:36:44 -0700 Subject: [PATCH 2/6] Pass the kernels test's failure function as it is Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-kernels/tests/kernels_vs_sql.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/einfold-kernels/tests/kernels_vs_sql.rs b/crates/einfold-kernels/tests/kernels_vs_sql.rs index 67a965ee..226e9a87 100644 --- a/crates/einfold-kernels/tests/kernels_vs_sql.rs +++ b/crates/einfold-kernels/tests/kernels_vs_sql.rs @@ -746,7 +746,7 @@ async fn fold_matches_sql() { dense_ran += usize::from(ran == Algorithm::Dense); if ran == Algorithm::Dense && case.domain.is_some() { unique_checked += 1; - unique_accepted += usize::from(check_unique_domain(&case, &mut rng, |w| fail(w))); + unique_accepted += usize::from(check_unique_domain(&case, &mut rng, fail)); } let (ours, again) = (rows(&ours, keys), rows(&again, keys)); // SQL doesn't order rows, so the outputs are compared as multisets. From c481188c7540932ed747ccbc20ffbc6158d77e25 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 26 Oct 2026 11:36:44 -0700 Subject: [PATCH 3/6] Detect a fold joined back to its domain, and check its keys when it runs ddx writes each gradient step as a fold joined back to the table being differentiated, grouped again by that table's keys, then left-joined from it. When no two rows of the table reach one of the fold's groups, that is the fold with the table as its domain. DetectFolds now finds the shape, and the node runs the dense algorithm with a check that the domain's keys are unique, falling back to the host's plan when they aren't. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/src/exec.rs | 25 +- crates/einfold-datafusion/src/join_back.rs | 341 ++++++++++++++++++ crates/einfold-datafusion/src/lib.rs | 1 + crates/einfold-datafusion/src/node.rs | 45 ++- crates/einfold-datafusion/src/rule.rs | 20 +- .../tests/detect_vs_host.rs | 165 +++++++-- 6 files changed, 552 insertions(+), 45 deletions(-) create mode 100644 crates/einfold-datafusion/src/join_back.rs diff --git a/crates/einfold-datafusion/src/exec.rs b/crates/einfold-datafusion/src/exec.rs index 199d5076..b8ee1d17 100644 --- a/crates/einfold-datafusion/src/exec.rs +++ b/crates/einfold-datafusion/src/exec.rs @@ -70,6 +70,8 @@ pub struct EinFoldExec { inputs: Vec>, /// The host's plan for the same rows, run instead of the hash algorithm. fallback: Option>, + /// Whether to run the fallback unless no two domain rows reach one group. + unique_domain: bool, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -88,6 +90,7 @@ impl EinFoldExec { fold, inputs, fallback: None, + unique_domain: false, properties: Arc::new(properties), metrics: ExecutionPlanMetricsSet::new(), } @@ -101,6 +104,14 @@ impl EinFoldExec { self } + /// This operator, computing its fold only when no two of its domain's + /// rows reach one group, and running its fallback otherwise. It needs a + /// fallback (see [`EinFoldExec::with_fallback`]). + pub fn requiring_unique_domain(mut self) -> Self { + self.unique_domain = true; + self + } + /// The host's plan this operator runs where the dense algorithm doesn't /// apply, if it has one. pub fn fallback(&self) -> Option<&Arc> { @@ -111,6 +122,7 @@ impl EinFoldExec { fn copy(&self, inputs: Vec>) -> EinFoldExec { let mut exec = EinFoldExec::new(self.fold.clone(), inputs, self.schema()); exec.fallback = self.fallback.clone(); + exec.unique_domain = self.unique_domain; exec } @@ -122,7 +134,11 @@ impl EinFoldExec { impl DisplayAs for EinFoldExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "EinFoldExec: {}", describe(&self.fold)) + write!( + f, + "EinFoldExec: {}", + describe(&self.fold, self.unique_domain) + ) } } @@ -234,7 +250,11 @@ impl ExecutionPlan for EinFoldExec { let options = Options { host: HOST, threads: context.session_config().target_partitions(), + unique_domain: self.unique_domain, }; + if self.unique_domain && self.fallback.is_none() { + return internal_err!("EinFoldExec requires a unique domain, but has no fallback"); + } let builder = || MetricBuilder::new(&self.metrics); let dense_folds = builder().global_counter("dense_folds"); let hash_folds = builder().global_counter("hash_folds"); @@ -357,6 +377,9 @@ impl ExtensionPlanner for EinFoldPlanner { if let Some(plan) = node.fallback() { exec = exec.with_fallback(planner.create_physical_plan(plan, session).await?); } + if node.unique_domain() { + exec = exec.requiring_unique_domain(); + } Ok(Some(Arc::new(exec))) } } diff --git a/crates/einfold-datafusion/src/join_back.rs b/crates/einfold-datafusion/src/join_back.rs new file mode 100644 index 00000000..965c5b0a --- /dev/null +++ b/crates/einfold-datafusion/src/join_back.rs @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: 2026 Alexander Merose & einfold Authors +// +// SPDX-License-Identifier: Apache-2.0 + +//! The third shape [`DetectFolds`](crate::DetectFolds) finds: a fold joined +//! back to its domain. In SQL, +//! +//! ```sql +//! SELECT +//! FROM d LEFT JOIN ( +//! SELECT d2., SUM() AS s +//! FROM d2 JOIN f ON d2. = f. AND … +//! GROUP BY d2. +//! ) r ON d. = r. AND … +//! ``` +//! +//! where `f` is a fold the rule found, `d2` is `d` under other names, the +//! inner join matches every one of `f`'s outputs, the second `GROUP BY` is by +//! exactly those keys, and `e` is `f`'s value, or `NULL` where a test of +//! `d2`'s row holds. ddx, an XQL Systems project for automatic +//! differentiation of SQL queries, writes each gradient step this way. +//! +//! When no two rows of `d` have the keys of one of `f`'s groups, each group +//! of the second `GROUP BY` holds one row, and the query is the fold with `d` +//! as its domain: each row of `d` gets `0.0 + e` evaluated on itself and its +//! group's value, where `0.0 +` is what `SUM` of one value computes, and NULL +//! where no group matched. `e` is NULL where `f`'s value is, so a missing +//! group and a group with a NULL sum agree. Whether the keys are unique is +//! known only when the query runs, so the node checks it then, and runs the +//! host's own plan where they aren't. + +use std::sync::Arc; + +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{Column, NullEquality, Result}; +use datafusion::logical_expr::expr::Case; +use datafusion::logical_expr::{lit, Expr, Join, JoinType, LogicalPlan, Projection}; +use einfold_ir::{ColumnIndex, Domain, Fold, KeyEquality, Output}; + +use crate::node::FoldNode; +use crate::rule::{found_fold, resolve, restore, same_schema, sum_argument, PREFIX}; + +/// `plan` without the aliases and column-only projections over it, with, for +/// each of `plan`'s columns, its column in what is left. A table scan's +/// columns are given as the table's own, whichever of them it reads. +fn base(plan: &LogicalPlan) -> Option<(&LogicalPlan, Vec)> { + match plan { + LogicalPlan::SubqueryAlias(a) => base(&a.input), + LogicalPlan::Projection(p) => { + let (b, below) = base(&p.input)?; + let columns = (p.expr.iter()) + .map(|e| match e.clone().unalias() { + Expr::Column(c) => p.input.schema().index_of_column(&c).ok().map(|i| below[i]), + _ => None, + }) + .collect::>>()?; + Some((b, columns)) + } + LogicalPlan::TableScan(scan) => match &scan.projection { + Some(columns) => Some((plan, columns.clone())), + None => Some((plan, (0..scan.source.schema().fields().len()).collect())), + }, + other => Some((other, (0..other.schema().fields().len()).collect())), + } +} + +/// Whether two bases (see [`base`]) are one relation: the same plan, or +/// scans of one table with the same filters and limit, whichever columns +/// each reads. +fn same_relation(a: &LogicalPlan, b: &LogicalPlan) -> bool { + match (a, b) { + (LogicalPlan::TableScan(a), LogicalPlan::TableScan(b)) => { + a.table_name == b.table_name + && Arc::ptr_eq(&a.source, &b.source) + && a.filters == b.filters + && a.fetch == b.fetch + } + (a, b) => a == b, + } +} + +/// Whether `e` is `value`, or `CASE WHEN THEN NULL ELSE value END` +/// with `test` reading only `row`'s columns: NULL wherever `value` is. +fn null_with_value(e: &Expr, value: &Column, row: &LogicalPlan) -> bool { + let is_value = |e: &Expr| matches!(e, Expr::Column(c) if c == value); + match e { + Expr::Case(Case { + expr: None, + when_then_expr, + else_expr: Some(otherwise), + }) => { + is_value(otherwise) + && when_then_expr.iter().all(|(test, then)| { + matches!(then.as_ref(), Expr::Literal(v, _) if v.is_null()) + && (test.column_refs().iter()).all(|c| row.schema().has_column(c)) + }) + } + e => is_value(e), + } +} + +/// `p` over a fold joined back to its domain (see the module's +/// documentation), as the fold from its domain, if that is what `p` is. +pub(crate) fn joined_back(p: &Projection) -> Result> { + let LogicalPlan::Join(Join { + left: d, + right, + on, + filter: None, + join_type: JoinType::Left, + null_equality: outer, + .. + }) = p.input.as_ref() + else { + return Ok(None); + }; + // The right side: an aggregate with one SUM, under aliases and renames; + // each of its columns is one of the aggregate's. + let Some((LogicalPlan::Aggregate(agg), right_map)) = base(right) else { + return Ok(None); + }; + let Some(arg) = sum_argument(agg) else { + return Ok(None); + }; + // Down through projections to the inner join of d2 and the fold. + let mut projections = Vec::new(); + let mut below = agg.input.as_ref(); + let inner = loop { + match below { + LogicalPlan::Projection(q) => { + projections.push(q); + below = q.input.as_ref(); + } + LogicalPlan::Join(j) => break j, + _ => return Ok(None), + } + }; + if inner.join_type != JoinType::Inner || inner.filter.is_some() { + return Ok(None); + } + let (d2, f, (node, f_columns), flip) = match (found_fold(&inner.left), found_fold(&inner.right)) + { + (None, Some(found)) => (&inner.left, &inner.right, found, false), + (Some(found), None) => (&inner.right, &inner.left, found, true), + _ => return Ok(None), + }; + let fold = node.fold(); + let n = fold.outputs().len(); + // d and d2 must be one relation under two sets of names. + let (Some((d_base, d_map)), Some((d2_base, d2_map))) = (base(d), base(d2)) else { + return Ok(None); + }; + if !same_relation(d_base, d2_base) { + return Ok(None); + } + // The inner join pairs a column of d2 with each of the fold's outputs. + let mut output_key: Vec> = vec![None; n]; + for (l, r) in &inner.on { + let (dk, fk) = if flip { (r, l) } else { (l, r) }; + let (Expr::Column(dc), Expr::Column(fc)) = (dk, fk) else { + return Ok(None); + }; + let (Ok(di), Ok(fi)) = ( + d2.schema().index_of_column(dc), + f.schema().index_of_column(fc), + ) else { + return Ok(None); + }; + let o = f_columns[fi]; + if o >= n || output_key[o].replace(d2_map[di]).is_some() { + return Ok(None); + } + } + let Some(output_key) = output_key.into_iter().collect::>>() else { + return Ok(None); + }; + // The second GROUP BY is by exactly those columns, each once. + let mut grouped = Vec::new(); + for g in &agg.group_expr { + let Expr::Column(c) = resolve(g, &projections)? else { + return Ok(None); + }; + let Ok(i) = d2.schema().index_of_column(&c) else { + return Ok(None); + }; + grouped.push(d2_map[i]); + } + let (mut a, mut b) = (grouped.clone(), output_key.clone()); + a.sort(); + b.sort(); + a.dedup(); + if a != b || grouped.len() != n { + return Ok(None); + } + // The left join pairs each group key with its own column of d. + let right_schema = right.schema(); + let mut joined = vec![false; n]; + for (l, r) in on { + let (Expr::Column(lc), Expr::Column(rc)) = (l, r) else { + return Ok(None); + }; + let (Ok(di), Ok(j)) = ( + d.schema().index_of_column(lc), + right_schema.index_of_column(rc), + ) else { + return Ok(None); + }; + let j = right_map[j]; + if j >= n || grouped[j] != d_map[di] || std::mem::replace(&mut joined[j], true) { + return Ok(None); + } + } + if joined.contains(&false) { + return Ok(None); + } + // A NULL key the inner join matched, and `=` in the left join didn't. + if *outer == NullEquality::NullEqualsNothing + && inner.null_equality == NullEquality::NullEqualsNull + { + return Ok(None); + } + // The sum's argument: the fold's value, NULL wherever it is. + let Some(v) = f_columns.iter().position(|&c| c == n) else { + return Ok(None); + }; + let value = Column::from(f.schema().qualified_field(v)); + let e = resolve(&arg, &projections)?; + if !null_with_value(&e, &value, d2) { + return Ok(None); + } + // p may read d's columns and the sum, not the second GROUP BY's keys. + let Some(s) = right_map.iter().position(|&c| c == n) else { + return Ok(None); + }; + let sum = Column::from(right_schema.qualified_field(s)); + let reads_keys = p.expr.iter().any(|e| { + (e.column_refs().iter()).any(|c| { + let i = right_schema.index_of_column(c).ok(); + i.is_some_and(|i| right_map[i] != n) + }) + }); + if reads_keys { + return Ok(None); + } + + // The fold, with d as its domain: each output matched with its column + // of d, under the inner join's equality. + let d_column = |b: usize| d_map.iter().position(|&c| c == b); + let equality = match inner.null_equality { + NullEquality::NullEqualsNull => KeyEquality::NotDistinct, + NullEquality::NullEqualsNothing => KeyEquality::Equal, + }; + let mut outputs = Vec::new(); + for (o, &b) in fold.outputs().iter().zip(&output_key) { + let Some(c) = d_column(b) else { + return Ok(None); + }; + outputs.push(Output { + domain_key: Some((ColumnIndex(c), equality)), + ..*o + }); + } + let width = d.schema().fields().len(); + let Ok(with_domain) = Fold::try_new( + fold.dimensions().to_vec(), + fold.operands().to_vec(), + outputs, + fold.row(), + fold.aggregate().clone(), + Some(Domain { + carry: (0..width).map(ColumnIndex).collect(), + missing: None, + }), + ) else { + return Ok(None); + }; + // The plan it replaces: the left join, as d's columns and the sum. + let mut kept: Vec = d.schema().columns().into_iter().map(Expr::Column).collect(); + kept.push(Expr::Column(sum.clone())); + let original = LogicalPlan::Projection(Projection::try_new(kept, Arc::clone(&p.input))?); + let names: Vec = (d.schema().fields().iter()) + .map(|f| f.name().clone()) + .chain([format!("{PREFIX}v")]) + .collect(); + let inputs = node.inputs_owned(); + let Ok(new) = FoldNode::try_new( + with_domain, + inputs[..2].to_vec(), + Some(d.as_ref().clone()), + names.clone(), + ) else { + return Ok(None); + }; + let new = new + .with_fallback(original.clone()) + .requiring_unique_domain() + .into_plan(); + // Over it: d's columns, and the sum as SUM of one value gives it, `e` + // reading d2's columns from the node's copy of d. + let mut unmapped = false; + let renamed = e + .transform(|x| match &x { + Expr::Column(c) if *c == value => Ok(Transformed::yes(Expr::Column( + Column::new_unqualified(&names[width]), + ))), + Expr::Column(c) => { + let column = d2.schema().index_of_column(c).ok(); + match column.and_then(|i| d_column(d2_map[i])) { + Some(i) => Ok(Transformed::yes(Expr::Column(Column::new_unqualified( + &names[i], + )))), + None => { + unmapped = true; + Ok(Transformed::no(x)) + } + } + } + _ => Ok(Transformed::no(x)), + })? + .data; + if unmapped { + return Ok(None); + } + let Some(columns) = restore(new, &names[..width], d.schema())? else { + return Ok(None); + }; + let LogicalPlan::Projection(columns) = columns else { + return Ok(None); + }; + let mut exprs = columns.expr.clone(); + exprs.push((lit(0.0) + renamed).alias_qualified(sum.relation.clone(), &sum.name)); + let projection = Projection::try_new(exprs, Arc::clone(&columns.input))?; + if !same_schema(&projection.schema, original.schema()) { + return Ok(None); + } + let projection = Projection::try_new( + p.expr.clone(), + Arc::new(LogicalPlan::Projection(projection)), + )?; + Ok(same_schema(&projection.schema, &p.schema).then_some(LogicalPlan::Projection(projection))) +} diff --git a/crates/einfold-datafusion/src/lib.rs b/crates/einfold-datafusion/src/lib.rs index 0cc102a9..760c440c 100644 --- a/crates/einfold-datafusion/src/lib.rs +++ b/crates/einfold-datafusion/src/lib.rs @@ -16,6 +16,7 @@ //! engine. mod exec; +mod join_back; mod node; mod rule; diff --git a/crates/einfold-datafusion/src/node.rs b/crates/einfold-datafusion/src/node.rs index 326badba..a9c56c08 100644 --- a/crates/einfold-datafusion/src/node.rs +++ b/crates/einfold-datafusion/src/node.rs @@ -30,6 +30,9 @@ pub struct FoldNode { fallback: Option, /// Whether [`DetectFolds`](crate::DetectFolds) made this node. found: bool, + /// Whether the node computes its fold only when no two domain rows reach + /// one group, and runs its fallback otherwise. + unique_domain: bool, } impl FoldNode { @@ -71,6 +74,7 @@ impl FoldNode { schema, fallback: None, found: false, + unique_domain: false, }) } @@ -84,6 +88,21 @@ impl FoldNode { self } + /// This node, computing its fold only when no two of its domain's rows + /// reach one group, and running its fallback otherwise: for a plan that + /// joins a fold back to its domain and groups again by the domain's keys, + /// which is the fold from its domain exactly then. + pub(crate) fn requiring_unique_domain(mut self) -> Self { + self.unique_domain = true; + self + } + + /// Whether the node computes its fold only when its domain's keys are + /// unique, running its fallback otherwise. + pub fn unique_domain(&self) -> bool { + self.unique_domain + } + /// This node, marked as made by [`DetectFolds`](crate::DetectFolds) /// without a fallback. pub(crate) fn marked_found(mut self) -> Self { @@ -122,12 +141,10 @@ impl FoldNode { // The schema follows from the fold and the inputs, so ordering ignores it. impl PartialOrd for FoldNode { fn partial_cmp(&self, other: &Self) -> Option { - (&self.fold, &self.inputs, &self.fallback, self.found).partial_cmp(&( - &other.fold, - &other.inputs, - &other.fallback, - other.found, - )) + fn key(n: &FoldNode) -> impl PartialOrd + '_ { + (&n.fold, &n.inputs, &n.fallback, n.found, n.unique_domain) + } + key(self).partial_cmp(&key(other)) } } @@ -149,7 +166,7 @@ impl UserDefinedLogicalNodeCore for FoldNode { } fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "Fold: {}", describe(&self.fold)) + write!(f, "Fold: {}", describe(&self.fold, self.unique_domain)) } fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { @@ -167,6 +184,7 @@ impl UserDefinedLogicalNodeCore for FoldNode { Ok(FoldNode { fallback: self.fallback.clone(), found: self.found, + unique_domain: self.unique_domain, ..node }) } @@ -174,9 +192,10 @@ impl UserDefinedLogicalNodeCore for FoldNode { /// The fold in einsum notation for `EXPLAIN`: each operand's dimensions, /// then the outputs, such as `a,b * b,c -> a,c` for a matrix product, with -/// `+ domain` when there is one. Dimensions are lettered by number, as the +/// `+ domain` when there is one (`+ unique domain` when the node requires +/// its keys to be unique). Dimensions are lettered by number, as the /// fold names none; past `z`, they are shown as `#26`, `#27`, …. -pub(crate) fn describe(fold: &Fold) -> String { +pub(crate) fn describe(fold: &Fold, unique_domain: bool) -> String { let name = |DimensionId(d): DimensionId| match u8::try_from(d) { Ok(d) if d < 26 => char::from(b'a' + d).to_string(), _ => format!("#{d}"), @@ -197,10 +216,10 @@ pub(crate) fn describe(fold: &Fold) -> String { "{} -> {}{}", operands.join(" * "), outputs.join(","), - if fold.domain().is_some() { - " + domain" - } else { - "" + match (fold.domain(), unique_domain) { + (None, _) => "", + (Some(_), false) => " + domain", + (Some(_), true) => " + unique domain", } ) } diff --git a/crates/einfold-datafusion/src/rule.rs b/crates/einfold-datafusion/src/rule.rs index 8e561bbf..6b931f3b 100644 --- a/crates/einfold-datafusion/src/rule.rs +++ b/crates/einfold-datafusion/src/rule.rs @@ -48,11 +48,12 @@ use einfold_ir::{ Operand, Output, RowValue, }; +use crate::join_back::joined_back; use crate::node::FoldNode; /// The prefix of the names this rule gives the columns of the nodes it makes, /// chosen not to clash with a domain's own columns. -const PREFIX: &str = "__einfold_"; +pub(crate) const PREFIX: &str = "__einfold_"; /// Finds folds over joins in DataFusion's optimized plans (see the module's /// documentation). [`EinFoldQueryPlanner::detecting`](crate::EinFoldQueryPlanner::detecting) @@ -96,7 +97,12 @@ impl OptimizerRule for DetectFolds { ) -> Result> { let found = match &plan { LogicalPlan::Aggregate(agg) => fold(agg, self.fallback)?, - LogicalPlan::Projection(p) => from_domain(p, self.fallback)?, + LogicalPlan::Projection(p) => match from_domain(p, self.fallback)? { + // A fold joined back to its domain needs the host's plan + // for when the domain's keys aren't unique. + None if self.fallback => joined_back(p)?, + found => found, + }, _ => None, }; Ok(match found { @@ -110,7 +116,7 @@ impl OptimizerRule for DetectFolds { /// `e` with each column a projection in `projections` (outermost first) /// produced replaced by the expression that produced it, aliases removed. -fn resolve(e: &Expr, projections: &[&Projection]) -> Result { +pub(crate) fn resolve(e: &Expr, projections: &[&Projection]) -> Result { let mut e = e.clone().unalias(); for p in projections { e = e @@ -225,7 +231,7 @@ fn widens(from: &DataType, to: &DataType) -> bool { } /// The sum's argument, if `agg` computes exactly one plain `SUM`. -fn sum_argument(agg: &Aggregate) -> Option { +pub(crate) fn sum_argument(agg: &Aggregate) -> Option { let [aggregate] = agg.aggr_expr.as_slice() else { return None; }; @@ -423,7 +429,7 @@ pub(crate) fn host_plan(plan: LogicalPlan) -> Result { /// `node`, whose columns are `names`, renamed to `schema`'s columns, or /// `None` if the result's schema wouldn't be exactly `schema`. -fn restore( +pub(crate) fn restore( node: LogicalPlan, names: &[String], schema: &DFSchemaRef, @@ -441,7 +447,7 @@ fn restore( /// Whether two schemas have the same columns, qualifiers, types and /// nullability, which an optimizer rule must keep. -fn same_schema(a: &DFSchemaRef, b: &DFSchemaRef) -> bool { +pub(crate) fn same_schema(a: &DFSchemaRef, b: &DFSchemaRef) -> bool { a.fields().len() == b.fields().len() && a.iter().zip(b.iter()).all(|((qa, fa), (qb, fb))| { qa == qb @@ -457,7 +463,7 @@ fn same_schema(a: &DFSchemaRef, b: &DFSchemaRef) -> bool { /// that only rename or reorder columns, with, for each column of `plan`, its /// column in the fold's output: an output's index, or the number of outputs /// for the sum. -fn found_fold(plan: &LogicalPlan) -> Option<(&FoldNode, Vec)> { +pub(crate) fn found_fold(plan: &LogicalPlan) -> Option<(&FoldNode, Vec)> { match plan { LogicalPlan::Extension(e) => { let node = e.node.as_any().downcast_ref::()?; diff --git a/crates/einfold-datafusion/tests/detect_vs_host.rs b/crates/einfold-datafusion/tests/detect_vs_host.rs index 45ad32f5..2c5bf694 100644 --- a/crates/einfold-datafusion/tests/detect_vs_host.rs +++ b/crates/einfold-datafusion/tests/detect_vs_host.rs @@ -20,13 +20,20 @@ //! `CASE`), which the rule must turn into a fold with `d` as its domain; //! - **look-alikes** it must leave alone: a mask (`AND l.a >= r.b`), no //! `GROUP BY`, `AVG` rather than `SUM`, and a sum that isn't a product; -//! and look-alikes of a fold from a domain, in which it must find the fold but give it no -//! domain: a projection that reads the fold's keys, and inner and right -//! joins. +//! and look-alikes of a fold from a domain, in which it must find the fold +//! but give it no domain: a projection that reads the fold's keys, and +//! inner and right joins; +//! - **folds joined back to their domain** (ddx's gradient step: the fold +//! joined to `d`, grouped again by `d`'s keys, and left-joined from `d`), +//! which the rule must turn into a fold from `d` that checks `d`'s keys are +//! unique when it runs, and their look-alikes, which keep the fold alone. //! //! Invariants: the same rows, as a multiset, bit for bit (every NaN alike), as -//! without einfold; the rule fires on every fold and fold from a domain, and never on a -//! look-alike. Half the cases run every fold the rule finds through einfold's +//! without einfold; the rule fires on every fold, fold from a domain and fold +//! joined back, and never on a look-alike. A second test draws only folds +//! joined back, over operands the dense algorithm takes and a domain whose +//! keys are duplicated half the time, so that both outcomes of the +//! uniqueness check are exercised. Half the cases run every fold the rule finds through einfold's //! own kernels (`DetectFolds::without_fallback`), so that the folds the rule //! describes are checked even where the dense algorithm doesn't apply. //! @@ -164,12 +171,17 @@ enum Expect { FromDomain, /// A look-alike of a fold from a domain: the fold is found, but it gets no domain. FoldOnly, + /// A fold joined back to its domain and grouped again, as ddx writes a + /// gradient: with the host's plan as a fallback, a fold from a domain + /// whose keys must be unique; without one, the fold alone. + JoinedBack, /// No fold at all. Decline, } -/// A random query over `l`, `r` and `d`, and whether the rule must fire. -fn query(rng: &mut Rng) -> (String, Expect) { +/// A random query over `l`, `r` and `d`, and whether the rule must fire; +/// with `joined_back`, always a fold joined back to `d` or a look-alike. +fn query(rng: &mut Rng, joined_back: bool) -> (String, Expect) { let eq = |x: &str, y: &str, safe: bool| match safe { true => format!("({x} IS NOT DISTINCT FROM {y})"), false => format!("{x} = {y}"), @@ -177,7 +189,7 @@ fn query(rng: &mut Rng) -> (String, Expect) { let safe = rng.chance(0.5); // Join on l.b = r.a, sometimes also l.a = r.b. let mut on = vec![eq("l.b", "r.a", safe)]; - let two_keys = rng.chance(0.3); + let two_keys = rng.chance(0.3) && !joined_back; if two_keys { on.push(eq("l.a", "r.b", safe)); } @@ -189,6 +201,7 @@ fn query(rng: &mut Rng) -> (String, Expect) { let (fl, fr) = (factor("l", rng), factor("r", rng)); // Group keys: l.a and r.b, or one of them, or the join key. let groups: Vec<&str> = match rng.below(4) { + 0 | 1 | 3 if joined_back => vec!["l.a", "r.b"], 0 => vec!["l.a", "r.b"], 1 => vec!["l.a"], 2 => vec!["r.b", "l.a"], @@ -215,8 +228,68 @@ fn query(rng: &mut Rng) -> (String, Expect) { ) }; let product = format!("SUM({fl} * {fr})"); - match rng.below(8) { + let shape = rng.below(10); + match if joined_back { 8 } else { shape } { 0..=2 => (fold(&product, false), Expect::Fold), + 8..=9 if groups.len() == 2 => { + // ddx's gradient step: the fold joined back to d, grouped again by + // d's keys, then left-joined from d. Its look-alikes keep the fold + // alone: a second GROUP BY by other columns, a sum that isn't NULL where + // the fold's value is, a filtered d, a projection reading the + // second GROUP BY's keys, and `=` outside `IS NOT DISTINCT FROM`. + let (inner, outer) = match rng.below(8) { + 0 => (true, false), + _ => { + let s = rng.chance(0.7); + (s, s) + } + }; + let look = rng.below(11); + // The second GROUP BY's second column: the join's, or d's value. + let col = if look == 10 { "v" } else { "b" }; + let inner_keys = [eq("d2.a", "f.g0", inner), eq("d2.b", "f.g1", inner)].join(" AND "); + let outer_keys = [ + eq("d.a", "s.a", outer), + eq(&format!("d.{col}"), "s.b", outer), + ]; + let outer_keys = outer_keys.join(" AND "); + let (group, sum, d2, reads) = match look { + 0 => ("d2.a, d2.b", "SUM(f.v * 2.0)", "d", "d.b"), + 1 => ("d2.a, d2.b", "SUM(COALESCE(f.v, 1.0))", "d", "d.b"), + 2 => ( + "d2.a, d2.b", + "SUM(f.v)", + "(SELECT * FROM d WHERE d.v IS NOT NULL)", + "d.b", + ), + 3 => ("d2.a, d2.b", "SUM(f.v)", "d", "s.b"), + 4 => ("d2.b, d2.a", "SUM(f.v)", "d", "d.b"), + 5 => ("d2.a, d2.b", "SUM(f.v)", "(SELECT a, b, v FROM d)", "d.b"), + 10 => ("d2.a, d2.v", "SUM(f.v)", "d", "d.b"), + _ => ( + "d2.a, d2.b", + "SUM(CASE WHEN d2.v IS NULL THEN NULL ELSE f.v END)", + "d", + "d.b", + ), + }; + let expect = match (look, inner, outer) { + (0..=3 | 10, _, _) | (_, true, false) => Expect::FoldOnly, + _ => Expect::JoinedBack, + }; + let regroup = format!( + "SELECT d2.a AS a, d2.{col} AS b, {sum} AS v FROM {d2} d2 JOIN ({}) f ON {inner_keys} \ + GROUP BY {group}", + fold(&product, false) + ); + ( + format!( + "SELECT d.a, {reads}, CASE WHEN d.v IS NULL THEN NULL WHEN s.v IS NULL THEN 0.0 \ + ELSE s.v END AS v FROM d LEFT JOIN ({regroup}) s ON {outer_keys}" + ), + expect, + ) + } 3..=4 if groups.len() == 2 => { // A fold from a domain, as ddx writes it: every row of d gets its group's sum, 0 where none. // Its look-alikes, which keep the fold but must not become a @@ -243,7 +316,7 @@ fn query(rng: &mut Rng) -> (String, Expect) { expect, ) } - 3..=4 => (fold(&product, false), Expect::Fold), + 3..=4 | 8..=9 => (fold(&product, false), Expect::Fold), 5 => (fold(&product, true), Expect::Decline), 6 => ( [ @@ -341,16 +414,50 @@ fn rows(batches: &[RecordBatch]) -> Vec { #[tokio::test] async fn einfold_changes_no_answer() { + let (cases, dense, host_plans, joined) = sweep(0xde7e_c7ed, false).await; + // Most spellings rule the dense path out (a group key that is a join key, + // a summed key held by one side, keys that can't be positions), and the + // kernels' own test covers it. Here both it and the host's plan must run. + if cases >= 100 { + let some = |n: usize| n * 50 >= cases; + assert!(some(dense), "the dense path ran in only {dense} cases"); + assert!( + some(host_plans), + "the host's plan ran in only {host_plans} cases" + ); + assert!(joined.0 + joined.1 > 0, "no fold joined back ran"); + } +} + +#[tokio::test] +async fn folds_joined_back_check_their_domain() { + let (cases, _, _, (dense, host)) = sweep(0x10_1ced, true).await; + if cases >= 100 { + let some = |n: usize| n * 10 >= cases; + assert!( + some(dense) && some(host), + "folds joined back ran dense in {dense} cases and on the host in {host}" + ); + } +} + +/// Runs the cases from `default_seed`, or as `EINFOLD_SEED` and +/// `EINFOLD_CASES` say, panicking at the first that fails; with +/// `joined_back`, only folds joined back to their domain. Returns how many +/// cases ran, how many ran the dense path and the host's plan, and how many +/// folds joined back ran each. +async fn sweep(default_seed: u64, joined_back: bool) -> (usize, usize, usize, (usize, usize)) { let seed: u64 = std::env::var("EINFOLD_SEED") .ok() .and_then(|s| s.parse().ok()) - .unwrap_or(0xde7e_c7ed); + .unwrap_or(default_seed); let cases: u64 = match std::env::var("EINFOLD_CASES") { Ok(n) => n.parse().unwrap(), Err(_) if std::env::var("EINFOLD_SEED").is_ok() => 1, Err(_) => 300, }; let (mut fired, mut dense, mut host_plans) = (0, 0, 0); + let (mut joined_dense, mut joined_host) = (0, 0); for i in 0..cases { let seed = seed.wrapping_add(i); let mut rng = Rng::new(seed); @@ -358,14 +465,21 @@ async fn einfold_changes_no_answer() { // Dense keys are sometimes clean enough for the dense algorithm. let messy = !matches!(keys, Keys::Dense) || rng.chance(0.4); let fallback = rng.chance(0.5); + let (keys, messy, fallback) = match joined_back { + true => (Keys::Dense, false, true), + false => (keys, messy, fallback), + }; let (ours, theirs) = (einfold(1 + rng.below(6), fallback), host(1 + rng.below(6))); let ctxs = [&ours, &theirs]; let ext = |rng: &mut Rng| 1 + rng.below(5); let (ea, eb, ec) = (ext(&mut rng), ext(&mut rng), ext(&mut rng)); table(&mut rng, &ctxs, "l", [ea, eb], keys, messy); table(&mut rng, &ctxs, "r", [eb, ec], keys, messy); - table(&mut rng, &ctxs, "d", [ea + 1, ec + 1], keys, messy); - let (sql, expect) = query(&mut rng); + // d is messy on its own sometimes, so that a domain with duplicate + // keys meets operands the dense algorithm can take. + let d_messy = messy || rng.chance(if joined_back { 0.5 } else { 0.3 }); + table(&mut rng, &ctxs, "d", [ea + 1, ec + 1], keys, d_messy); + let (sql, expect) = query(&mut rng, joined_back); let fail = |what: String| panic!("seed {seed}: {what}\nSQL: {sql}"); let df = ours.sql(&sql).await.unwrap(); @@ -382,6 +496,7 @@ async fn einfold_changes_no_answer() { let ok = match expect { Expect::Fold | Expect::FoldOnly => found == 1 && domains == 0, Expect::FromDomain => found == 1 && domains == 1, + Expect::JoinedBack => found == 1 && domains == usize::from(fallback), Expect::Decline => found == 0, }; if !ok { @@ -401,6 +516,10 @@ async fn einfold_changes_no_answer() { .to_string(); dense += usize::from(shown.contains("dense_folds=1")); host_plans += usize::from(shown.contains("host_plans=1")); + if expect == Expect::JoinedBack && fallback { + joined_dense += usize::from(shown.contains("dense_folds=1")); + joined_host += usize::from(shown.contains("host_plans=1")); + } let want = theirs.sql(&sql).await.unwrap().collect().await.unwrap(); let (got, want) = (rows(&got), rows(&want)); if got != want { @@ -410,17 +529,15 @@ async fn einfold_changes_no_answer() { eprintln!( "the rule fired in {fired} of {cases} cases; the dense path ran in {dense}, the host's plan in {host_plans}" ); - // Most spellings rule the dense path out (a group key that is a join key, - // a summed key held by one side, keys that can't be positions), and the - // kernels' own test covers it. Here both it and the host's plan must run. - if cases >= 100 { - let some = |n: usize| n * 50 >= cases as usize; - assert!(some(dense), "the dense path ran in only {dense} cases"); - assert!( - some(host_plans), - "the host's plan ran in only {host_plans} cases" - ); - } + eprintln!( + "folds joined back to their domain: the dense path ran in {joined_dense}, the host's plan in {joined_host}" + ); + ( + cases as usize, + dense, + host_plans, + (joined_dense, joined_host), + ) } /// Runs `setup`'s statements, then `sql` with and without einfold, and From 2af852a0f0eabd84aa413d8c26c36e7a33b8c778 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 26 Oct 2026 11:42:10 -0700 Subject: [PATCH 4/6] Match scans of one table by name, as ddx's plans give each scan its own handle Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/src/join_back.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/einfold-datafusion/src/join_back.rs b/crates/einfold-datafusion/src/join_back.rs index 965c5b0a..94469c54 100644 --- a/crates/einfold-datafusion/src/join_back.rs +++ b/crates/einfold-datafusion/src/join_back.rs @@ -66,12 +66,13 @@ fn base(plan: &LogicalPlan) -> Option<(&LogicalPlan, Vec)> { /// Whether two bases (see [`base`]) are one relation: the same plan, or /// scans of one table with the same filters and limit, whichever columns -/// each reads. +/// each reads. Within one plan, a table's name names one table, though +/// each scan of it may hold its own handle to it. fn same_relation(a: &LogicalPlan, b: &LogicalPlan) -> bool { match (a, b) { (LogicalPlan::TableScan(a), LogicalPlan::TableScan(b)) => { a.table_name == b.table_name - && Arc::ptr_eq(&a.source, &b.source) + && a.source.schema() == b.source.schema() && a.filters == b.filters && a.fetch == b.fetch } From 61154a1d8565d04126a73cc9c0fc8c9d1f5e6e51 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 26 Oct 2026 20:00:17 -0700 Subject: [PATCH 5/6] Apply a joined-back fold's tests in the node, not above it A fold joined back to its domain is detected as the fold from its domain, and its regroup's CASE (NULL where a test of the domain's row holds) was applied in a projection above the node. When the uniqueness check declined and the host's plan ran, the node's value was already the regroup's sum, so the tests were applied a second time: a row of a duplicated key that failed the test got NULL, where SQL gives every row of the group the same sum. The node now computes the tests as a Boolean column of its domain and makes the value NULL where one holds, only when it computes the fold itself. The tests then run on every row of the domain, so they must be unable to fail. Fixes #60. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/src/exec.rs | 49 +++++++++- crates/einfold-datafusion/src/join_back.rs | 98 +++++++++++++------ crates/einfold-datafusion/src/node.rs | 37 ++++++- crates/einfold-datafusion/src/rule.rs | 2 +- .../tests/detect_vs_host.rs | 51 ++++++++-- 5 files changed, 189 insertions(+), 48 deletions(-) diff --git a/crates/einfold-datafusion/src/exec.rs b/crates/einfold-datafusion/src/exec.rs index b8ee1d17..0b878ed4 100644 --- a/crates/einfold-datafusion/src/exec.rs +++ b/crates/einfold-datafusion/src/exec.rs @@ -9,8 +9,8 @@ use std::fmt; use std::sync::Arc; use async_trait::async_trait; -use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::compute::concat_batches; +use datafusion::arrow::array::{ArrayRef, BooleanArray, RecordBatch}; +use datafusion::arrow::compute::{concat_batches, nullif}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::catalog::Session; use datafusion::common::runtime::SpawnedTask; @@ -31,7 +31,7 @@ use datafusion::physical_plan::{ PlanProperties, }; use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; -use einfold_ir::Fold; +use einfold_ir::{ColumnIndex, Fold}; use einfold_kernels::{Algorithm, Options}; use futures::future::try_join_all; use futures::{stream, StreamExt, TryStreamExt}; @@ -72,6 +72,8 @@ pub struct EinFoldExec { fallback: Option>, /// Whether to run the fallback unless no two domain rows reach one group. unique_domain: bool, + /// The domain's Boolean column that makes a row's value NULL where true. + null_when: Option, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -91,6 +93,7 @@ impl EinFoldExec { inputs, fallback: None, unique_domain: false, + null_when: None, properties: Arc::new(properties), metrics: ExecutionPlanMetricsSet::new(), } @@ -112,6 +115,15 @@ impl EinFoldExec { self } + /// This operator, giving a NULL value to each row of its domain where the + /// domain's Boolean column `column` is true, as `CASE WHEN THEN + /// NULL ELSE END` does. The host's plan, if it runs, computes its + /// own values. + pub fn with_null_when(mut self, column: ColumnIndex) -> Self { + self.null_when = Some(column); + self + } + /// The host's plan this operator runs where the dense algorithm doesn't /// apply, if it has one. pub fn fallback(&self) -> Option<&Arc> { @@ -123,6 +135,7 @@ impl EinFoldExec { let mut exec = EinFoldExec::new(self.fold.clone(), inputs, self.schema()); exec.fallback = self.fallback.clone(); exec.unique_domain = self.unique_domain; + exec.null_when = self.null_when; exec } @@ -260,6 +273,7 @@ impl ExecutionPlan for EinFoldExec { let hash_folds = builder().global_counter("hash_folds"); let host_plans = builder().global_counter("host_plans"); let fallback = self.fallback.clone(); + let null_when = self.null_when; let compute = builder().elapsed_compute(0); let output_rows = builder().output_rows(0); // Shared with the kernel's task, which holds the inputs it counts: a @@ -279,6 +293,11 @@ impl ExecutionPlan for EinFoldExec { ) .await?; let domain = (batches.len() == 3).then(|| batches.pop().unwrap()); + // Each domain batch's NULL mask: the kernels give one output batch + // per domain batch, row for row. + let masks: Option> = null_when + .zip(domain.as_ref()) + .map(|(c, d)| d.iter().map(|b| Arc::clone(b.column(c.0))).collect()); let b = concat(&schemas[1], batches.pop().unwrap(), &reservation)?; let a = concat(&schemas[0], batches.pop().unwrap(), &reservation)?; let timer = compute.timer(); @@ -304,14 +323,31 @@ impl ExecutionPlan for EinFoldExec { // memory they would have had without einfold. The output, which // may share the domain's columns, is counted below. reservation.free(); + let masked = |batches: Vec| -> Result> { + let Some(masks) = &masks else { + return Ok(batches); + }; + (batches.iter().zip(masks)) + .map(|(b, mask)| { + let mask = mask.as_any().downcast_ref::(); + let Some(mask) = mask else { + return internal_err!("EinFoldExec's NULL mask isn't Boolean"); + }; + let mut columns = b.columns().to_vec(); + let last = columns.len() - 1; + columns[last] = nullif(&columns[last], mask)?; + Ok(RecordBatch::try_new(b.schema(), columns)?) + }) + .collect() + }; let batches = match (ran, fallback) { (Some((Algorithm::Dense, batches)), _) => { dense_folds.add(1); - batches + masked(batches)? } (Some((Algorithm::Hash, batches)), _) => { hash_folds.add(1); - batches + masked(batches)? } (None, Some(plan)) => { // The host's plan, its columns under this operator's names. @@ -380,6 +416,9 @@ impl ExtensionPlanner for EinFoldPlanner { if node.unique_domain() { exec = exec.requiring_unique_domain(); } + if let Some(c) = node.null_when() { + exec = exec.with_null_when(c); + } Ok(Some(Arc::new(exec))) } } diff --git a/crates/einfold-datafusion/src/join_back.rs b/crates/einfold-datafusion/src/join_back.rs index 94469c54..81707255 100644 --- a/crates/einfold-datafusion/src/join_back.rs +++ b/crates/einfold-datafusion/src/join_back.rs @@ -24,7 +24,9 @@ //! of the second `GROUP BY` holds one row, and the query is the fold with `d` //! as its domain: each row of `d` gets `0.0 + e` evaluated on itself and its //! group's value, where `0.0 +` is what `SUM` of one value computes, and NULL -//! where no group matched. `e` is NULL where `f`'s value is, so a missing +//! where no group matched. The node computes `e`'s tests as a column of its +//! domain, and makes the value NULL where one holds, so that the host's plan, +//! when it runs instead, isn't given them a second time. `e` is NULL where `f`'s value is, so a missing //! group and a group with a NULL sum agree. Whether the keys are unique is //! known only when the query runs, so the node checks it then, and runs the //! host's own plan where they aren't. @@ -38,7 +40,7 @@ use datafusion::logical_expr::{lit, Expr, Join, JoinType, LogicalPlan, Projectio use einfold_ir::{ColumnIndex, Domain, Fold, KeyEquality, Output}; use crate::node::FoldNode; -use crate::rule::{found_fold, resolve, restore, same_schema, sum_argument, PREFIX}; +use crate::rule::{found_fold, movable, resolve, restore, same_schema, sum_argument, PREFIX}; /// `plan` without the aliases and column-only projections over it, with, for /// each of `plan`'s columns, its column in what is left. A table scan's @@ -262,6 +264,48 @@ pub(crate) fn joined_back(p: &Projection) -> Result> { }); } let width = d.schema().fields().len(); + // `e`'s tests, as a column of d: true where `e` is NULL whatever the + // fold's value. The tests are then computed on every row of d, where the + // host's plan computes them only on rows that joined the fold, so they + // must be unable to fail. + let d_expr = |c: &Column| -> Option { + let i = d2.schema().index_of_column(c).ok()?; + let j = d_column(d2_map[i])?; + Some(Expr::Column(Column::from(d.schema().qualified_field(j)))) + }; + let null_when = match &e { + Expr::Case(Case { when_then_expr, .. }) => { + let mut unmapped = false; + let mut tests = Vec::new(); + for (test, _) in when_then_expr { + let test = (test.as_ref().clone()) + .transform(|x| match &x { + Expr::Column(c) => match d_expr(c) { + Some(e) => Ok(Transformed::yes(e)), + None => { + unmapped = true; + Ok(Transformed::no(x)) + } + }, + _ => Ok(Transformed::no(x)), + })? + .data; + // A CASE branch is taken where its test is true, not NULL. + tests.push(match test { + Expr::IsNull(_) | Expr::IsNotNull(_) => test, + test => Expr::IsTrue(Box::new(test)), + }); + } + let Some(mask) = tests.into_iter().reduce(Expr::or) else { + return Ok(None); + }; + if unmapped || !movable(&mask, d.schema()) { + return Ok(None); + } + Some(mask) + } + _ => None, + }; let Ok(with_domain) = Fold::try_new( fold.dimensions().to_vec(), fold.operands().to_vec(), @@ -284,44 +328,35 @@ pub(crate) fn joined_back(p: &Projection) -> Result> { .chain([format!("{PREFIX}v")]) .collect(); let inputs = node.inputs_owned(); + // The domain: d, and the mask after its columns. + let domain = match &null_when { + None => d.as_ref().clone(), + Some(mask) => { + let mut columns: Vec = + d.schema().columns().into_iter().map(Expr::Column).collect(); + columns.push(mask.clone().alias(format!("{PREFIX}null"))); + LogicalPlan::Projection(Projection::try_new(columns, Arc::clone(d))?) + } + }; let Ok(new) = FoldNode::try_new( with_domain, inputs[..2].to_vec(), - Some(d.as_ref().clone()), + Some(domain), names.clone(), ) else { return Ok(None); }; let new = new .with_fallback(original.clone()) - .requiring_unique_domain() - .into_plan(); - // Over it: d's columns, and the sum as SUM of one value gives it, `e` - // reading d2's columns from the node's copy of d. - let mut unmapped = false; - let renamed = e - .transform(|x| match &x { - Expr::Column(c) if *c == value => Ok(Transformed::yes(Expr::Column( - Column::new_unqualified(&names[width]), - ))), - Expr::Column(c) => { - let column = d2.schema().index_of_column(c).ok(); - match column.and_then(|i| d_column(d2_map[i])) { - Some(i) => Ok(Transformed::yes(Expr::Column(Column::new_unqualified( - &names[i], - )))), - None => { - unmapped = true; - Ok(Transformed::no(x)) - } - } - } - _ => Ok(Transformed::no(x)), - })? - .data; - if unmapped { - return Ok(None); + .requiring_unique_domain(); + let new = match null_when { + Some(_) => new.with_null_when(ColumnIndex(width)), + None => new, } + .into_plan(); + // Over it: d's columns, and the sum as SUM of one value gives it. The + // node has applied `e`'s tests where it computed the fold, and the host's + // plan has where it ran. let Some(columns) = restore(new, &names[..width], d.schema())? else { return Ok(None); }; @@ -329,7 +364,8 @@ pub(crate) fn joined_back(p: &Projection) -> Result> { return Ok(None); }; let mut exprs = columns.expr.clone(); - exprs.push((lit(0.0) + renamed).alias_qualified(sum.relation.clone(), &sum.name)); + let value = Expr::Column(Column::new_unqualified(&names[width])); + exprs.push((lit(0.0) + value).alias_qualified(sum.relation.clone(), &sum.name)); let projection = Projection::try_new(exprs, Arc::clone(&columns.input))?; if !same_schema(&projection.schema, original.schema()) { return Ok(None); diff --git a/crates/einfold-datafusion/src/node.rs b/crates/einfold-datafusion/src/node.rs index a9c56c08..fa1c11d0 100644 --- a/crates/einfold-datafusion/src/node.rs +++ b/crates/einfold-datafusion/src/node.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::common::{plan_err, DFSchema, DFSchemaRef, Result}; use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; -use einfold_ir::{DimensionId, Fold}; +use einfold_ir::{ColumnIndex, DimensionId, Fold}; /// A fold over a join in a DataFusion logical plan. Its inputs are the fold's /// two operands, then its domain if it has one; its schema is @@ -33,6 +33,9 @@ pub struct FoldNode { /// Whether the node computes its fold only when no two domain rows reach /// one group, and runs its fallback otherwise. unique_domain: bool, + /// A Boolean column of the domain whose rows get a NULL value where it is + /// true, whatever their group's. + null_when: Option, } impl FoldNode { @@ -75,6 +78,7 @@ impl FoldNode { fallback: None, found: false, unique_domain: false, + null_when: None, }) } @@ -103,6 +107,21 @@ impl FoldNode { self.unique_domain } + /// This node, giving a NULL value to each row of its domain where the + /// domain's Boolean column `column` is true, as `CASE WHEN THEN + /// NULL ELSE END` does. + pub(crate) fn with_null_when(mut self, column: ColumnIndex) -> Self { + self.null_when = Some(column); + self + } + + /// The domain's column that makes a row's value NULL where it is true, if + /// any: such a row gets a NULL value, as `CASE WHEN THEN NULL + /// ELSE END` gives it. + pub fn null_when(&self) -> Option { + self.null_when + } + /// This node, marked as made by [`DetectFolds`](crate::DetectFolds) /// without a fallback. pub(crate) fn marked_found(mut self) -> Self { @@ -142,7 +161,14 @@ impl FoldNode { impl PartialOrd for FoldNode { fn partial_cmp(&self, other: &Self) -> Option { fn key(n: &FoldNode) -> impl PartialOrd + '_ { - (&n.fold, &n.inputs, &n.fallback, n.found, n.unique_domain) + ( + &n.fold, + &n.inputs, + &n.fallback, + n.found, + n.unique_domain, + n.null_when, + ) } key(self).partial_cmp(&key(other)) } @@ -166,7 +192,11 @@ impl UserDefinedLogicalNodeCore for FoldNode { } fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "Fold: {}", describe(&self.fold, self.unique_domain)) + write!(f, "Fold: {}", describe(&self.fold, self.unique_domain))?; + match self.null_when { + Some(ColumnIndex(c)) => write!(f, ", NULL where domain column {c}"), + None => Ok(()), + } } fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { @@ -185,6 +215,7 @@ impl UserDefinedLogicalNodeCore for FoldNode { fallback: self.fallback.clone(), found: self.found, unique_domain: self.unique_domain, + null_when: self.null_when, ..node }) } diff --git a/crates/einfold-datafusion/src/rule.rs b/crates/einfold-datafusion/src/rule.rs index 6b931f3b..c8ee3bab 100644 --- a/crates/einfold-datafusion/src/rule.rs +++ b/crates/einfold-datafusion/src/rule.rs @@ -157,7 +157,7 @@ pub const MOVABLE_FUNCTIONS: &[&str] = &["abs", "exp", "ln", "tanh", "signum"]; /// only deterministic, row-wise expressions that can't raise an error move; /// a volatile one, such as `random()`, would also be drawn once per input /// row rather than once per joined pair. -fn movable(e: &Expr, schema: &DFSchema) -> bool { +pub(crate) fn movable(e: &Expr, schema: &DFSchema) -> bool { let ty = |e: &Expr| e.get_type(schema).ok(); let arithmetic = |e: &Expr| ty(e).is_some_and(|t| t.is_integer() || t.is_floating()); let fails = |e: &Expr| { diff --git a/crates/einfold-datafusion/tests/detect_vs_host.rs b/crates/einfold-datafusion/tests/detect_vs_host.rs index 2c5bf694..ef9795a1 100644 --- a/crates/einfold-datafusion/tests/detect_vs_host.rs +++ b/crates/einfold-datafusion/tests/detect_vs_host.rs @@ -253,6 +253,10 @@ fn query(rng: &mut Rng, joined_back: bool) -> (String, Expect) { eq(&format!("d.{col}"), "s.b", outer), ]; let outer_keys = outer_keys.join(" AND "); + // The regroup's test: ddx's, or one of another column, which + // duplicate rows of d may disagree on. + let test = ["d2.v IS NULL", "d2.v < 0.0", "d2.v IS NULL OR d2.v > 2.0"][rng.below(3)]; + let masked = format!("SUM(CASE WHEN {test} THEN NULL ELSE f.v END)"); let (group, sum, d2, reads) = match look { 0 => ("d2.a, d2.b", "SUM(f.v * 2.0)", "d", "d.b"), 1 => ("d2.a, d2.b", "SUM(COALESCE(f.v, 1.0))", "d", "d.b"), @@ -266,12 +270,7 @@ fn query(rng: &mut Rng, joined_back: bool) -> (String, Expect) { 4 => ("d2.b, d2.a", "SUM(f.v)", "d", "d.b"), 5 => ("d2.a, d2.b", "SUM(f.v)", "(SELECT a, b, v FROM d)", "d.b"), 10 => ("d2.a, d2.v", "SUM(f.v)", "d", "d.b"), - _ => ( - "d2.a, d2.b", - "SUM(CASE WHEN d2.v IS NULL THEN NULL ELSE f.v END)", - "d", - "d.b", - ), + _ => ("d2.a, d2.b", masked.as_str(), "d", "d.b"), }; let expect = match (look, inner, outer) { (0..=3 | 10, _, _) | (_, true, false) => Expect::FoldOnly, @@ -282,10 +281,16 @@ fn query(rng: &mut Rng, joined_back: bool) -> (String, Expect) { GROUP BY {group}", fold(&product, false) ); + // The value as ddx reads it, masked by d's own row, or plainly; + // d.v is read either way, so that the test's column is d's. + let value = [ + "CASE WHEN d.v IS NULL THEN NULL WHEN s.v IS NULL THEN 0.0 ELSE s.v END", + "s.v", + "COALESCE(s.v, 0.0)", + ][rng.below(3)]; ( format!( - "SELECT d.a, {reads}, CASE WHEN d.v IS NULL THEN NULL WHEN s.v IS NULL THEN 0.0 \ - ELSE s.v END AS v FROM d LEFT JOIN ({regroup}) s ON {outer_keys}" + "SELECT d.a, {reads}, d.v AS w, {value} AS v FROM d LEFT JOIN ({regroup}) s ON {outer_keys}" ), expect, ) @@ -741,3 +746,33 @@ async fn fallbacks_hold_only_the_hosts_plan() { execs(fallback, &mut inner); assert!(inner.is_empty(), "the fallback runs einfold again"); } + +/// A domain whose duplicate rows disagree on the regroup's test: the uniqueness +/// check declines, and the host's plan gives every row of a group the same +/// sum, which the node mustn't mask a second time (#60). +#[tokio::test] +async fn joined_back_tests_apply_once() { + let setup = [ + "CREATE TABLE l (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,1.0),(0,1,2.0),(1,0,3.0),(1,1,4.0)", + "CREATE TABLE r (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,1.0),(0,1,10.0),(1,0,100.0),(1,1,1000.0)", + "CREATE TABLE d (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,NULL),(0,0,5.0),(1,1,7.0)", + "CREATE TABLE e (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,NULL),(0,1,5.0),(1,1,-7.0)", + ]; + for (table, test) in [ + ("d", "d2.v IS NULL"), + ("d", "d2.v > 6.0"), + ("e", "d2.v < 0.0"), + ] { + let sql = format!( + "SELECT d.a, d.b, d.v AS w, COALESCE(s.v, 0.0) AS v FROM {table} d LEFT JOIN ( \ + SELECT d2.a AS a, d2.b AS b, SUM(CASE WHEN {test} THEN NULL ELSE f.v END) AS v \ + FROM {table} d2 JOIN (SELECT l.a AS g0, r.b AS g1, SUM(l.v * r.v) AS v \ + FROM l JOIN r ON l.b = r.a GROUP BY l.a, r.b) f \ + ON d2.a = f.g0 AND d2.b = f.g1 \ + GROUP BY d2.a, d2.b) s \ + ON d.a = s.a AND d.b = s.b" + ); + // Table e's keys are unique, so there the dense path runs. + assert!(same_answer(&setup, &sql).await > 0, "{sql}"); + } +} From e5fc00dce9cb8e1cdb214bbea45b62630994a2d7 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 27 Oct 2026 17:41:53 -0700 Subject: [PATCH 6/6] Fall back to the host's own plan for a fold joined back, without the inner fold node Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/src/join_back.rs | 9 +++- .../tests/detect_vs_host.rs | 47 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/crates/einfold-datafusion/src/join_back.rs b/crates/einfold-datafusion/src/join_back.rs index 81707255..6277c299 100644 --- a/crates/einfold-datafusion/src/join_back.rs +++ b/crates/einfold-datafusion/src/join_back.rs @@ -40,7 +40,9 @@ use datafusion::logical_expr::{lit, Expr, Join, JoinType, LogicalPlan, Projectio use einfold_ir::{ColumnIndex, Domain, Fold, KeyEquality, Output}; use crate::node::FoldNode; -use crate::rule::{found_fold, movable, resolve, restore, same_schema, sum_argument, PREFIX}; +use crate::rule::{ + found_fold, host_plan, movable, resolve, restore, same_schema, sum_argument, PREFIX, +}; /// `plan` without the aliases and column-only projections over it, with, for /// each of `plan`'s columns, its column in what is left. A table scan's @@ -322,7 +324,10 @@ pub(crate) fn joined_back(p: &Projection) -> Result> { // The plan it replaces: the left join, as d's columns and the sum. let mut kept: Vec = d.schema().columns().into_iter().map(Expr::Column).collect(); kept.push(Expr::Column(sum.clone())); - let original = LogicalPlan::Projection(Projection::try_new(kept, Arc::clone(&p.input))?); + // The host's own plan, with the inner fold swapped back for the plan it + // replaced, so that falling back doesn't run einfold again. + let input = host_plan(p.input.as_ref().clone())?; + let original = LogicalPlan::Projection(Projection::try_new(kept, Arc::new(input))?); let names: Vec = (d.schema().fields().iter()) .map(|f| f.name().clone()) .chain([format!("{PREFIX}v")]) diff --git a/crates/einfold-datafusion/tests/detect_vs_host.rs b/crates/einfold-datafusion/tests/detect_vs_host.rs index ef9795a1..1b2e9647 100644 --- a/crates/einfold-datafusion/tests/detect_vs_host.rs +++ b/crates/einfold-datafusion/tests/detect_vs_host.rs @@ -776,3 +776,50 @@ async fn joined_back_tests_apply_once() { assert!(same_answer(&setup, &sql).await > 0, "{sql}"); } } + +/// A fold joined back to its domain falls back to the host's own plan too. +#[tokio::test] +async fn joined_back_fallbacks_hold_only_the_hosts_plan() { + let ctx = einfold(4, true); + for s in [ + "CREATE TABLE l (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,1.0),(0,1,2.0),(1,0,3.0),(1,1,4.0)", + "CREATE TABLE r (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,1.0),(0,1,10.0),(1,0,100.0),(1,1,1000.0)", + "CREATE TABLE d (a BIGINT, b BIGINT, v DOUBLE) AS VALUES (0,0,NULL),(0,0,5.0),(1,1,7.0)", + ] { + ctx.sql(s).await.unwrap().collect().await.unwrap(); + } + let sql = "SELECT d.a, d.b, d.v AS w, COALESCE(s.v, 0.0) AS v FROM d LEFT JOIN ( \ + SELECT d2.a AS a, d2.b AS b, SUM(CASE WHEN d2.v IS NULL THEN NULL ELSE f.v END) AS v \ + FROM d d2 JOIN (SELECT l.a AS g0, r.b AS g1, SUM(l.v * r.v) AS v \ + FROM l JOIN r ON l.b = r.a GROUP BY l.a, r.b) f \ + ON d2.a = f.g0 AND d2.b = f.g1 \ + GROUP BY d2.a, d2.b) s \ + ON d.a = s.a AND d.b = s.b"; + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + fn execs( + plan: &Arc, + out: &mut Vec>, + ) { + let any: &dyn std::any::Any = plan.as_ref(); + if any.is::() { + out.push(Arc::clone(plan)); + } + for c in plan.children() { + execs(c, out); + } + } + let mut found = Vec::new(); + execs(&plan, &mut found); + assert_eq!(found.len(), 1, "one fold joined back"); + let any: &dyn std::any::Any = found[0].as_ref(); + let fallback = (any.downcast_ref::().unwrap().fallback()).expect("a fallback"); + let mut inner = Vec::new(); + execs(fallback, &mut inner); + assert!(inner.is_empty(), "the fallback runs einfold again"); +}