diff --git a/crates/einfold-datafusion/src/exec.rs b/crates/einfold-datafusion/src/exec.rs index 199d5076..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}; @@ -70,6 +70,10 @@ 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, + /// The domain's Boolean column that makes a row's value NULL where true. + null_when: Option, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -88,6 +92,8 @@ impl EinFoldExec { fold, inputs, fallback: None, + unique_domain: false, + null_when: None, properties: Arc::new(properties), metrics: ExecutionPlanMetricsSet::new(), } @@ -101,6 +107,23 @@ 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 + } + + /// 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> { @@ -111,6 +134,8 @@ 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.null_when = self.null_when; exec } @@ -122,7 +147,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,12 +263,17 @@ 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"); 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 @@ -259,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(); @@ -284,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. @@ -357,6 +413,12 @@ 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(); + } + 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 new file mode 100644 index 00000000..6277c299 --- /dev/null +++ b/crates/einfold-datafusion/src/join_back.rs @@ -0,0 +1,383 @@ +// 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. 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. + +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, 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 +/// 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. 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 + && a.source.schema() == b.source.schema() + && 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(); + // `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(), + 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())); + // 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")]) + .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(domain), + names.clone(), + ) else { + return Ok(None); + }; + let new = new + .with_fallback(original.clone()) + .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); + }; + let LogicalPlan::Projection(columns) = columns else { + return Ok(None); + }; + let mut exprs = columns.expr.clone(); + 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); + } + 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..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 @@ -30,6 +30,12 @@ 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, + /// 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 { @@ -71,6 +77,8 @@ impl FoldNode { schema, fallback: None, found: false, + unique_domain: false, + null_when: None, }) } @@ -84,6 +92,36 @@ 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, 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 { @@ -122,12 +160,17 @@ 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, + n.null_when, + ) + } + key(self).partial_cmp(&key(other)) } } @@ -149,7 +192,11 @@ 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))?; + 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 { @@ -167,6 +214,8 @@ impl UserDefinedLogicalNodeCore for FoldNode { Ok(FoldNode { fallback: self.fallback.clone(), found: self.found, + unique_domain: self.unique_domain, + null_when: self.null_when, ..node }) } @@ -174,9 +223,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 +247,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..c8ee3bab 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 @@ -151,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| { @@ -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..1b2e9647 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,73 @@ 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 "); + // 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"), + 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", masked.as_str(), "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) + ); + // 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}, d.v AS w, {value} 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 +321,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 +419,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 +470,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 +501,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 +521,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 +534,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 @@ -624,3 +746,80 @@ 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}"); + } +} + +/// 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"); +} 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..226e9a87 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, fail)); + } 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.