Run the host's plan where memory runs short, as its spilling operators would

crates/einfold-datafusion/src/detect_folds.rs

@@ -31,11 +31,13 @@
 //! below the join, and keeps every plan's schema exactly: names, qualifiers,
 //! types and nullability. [`EinFoldExec`](crate::EinFoldExec) runs each
 //! node by the dense algorithm where its keys and rows allow, and by the
-//! discrimination algorithm otherwise. A node that can be computed only under
-//! a condition known when it runs, such as a fold joined back to its domain,
-//! keeps the plan it replaced and runs that where the condition fails. With
-//! [`DetectFolds::with_host_fallback`], every node keeps it, and runs it
-//! wherever the dense algorithm doesn't apply.
+//! discrimination algorithm otherwise. Every node keeps the plan it
+//! replaced, the host's own, and runs it where the session's memory pool
+//! can't give einfold what it needs: the host's operators can spill to disk.
+//! A node that can be computed only under a condition known when it runs,
+//! such as a fold joined back to its domain, also runs it where the
+//! condition fails. With [`DetectFolds::with_host_fallback`], every node
+//! runs it wherever the dense algorithm doesn't apply.
 
 use std::any::Any;
 use std::sync::Arc;
@@ -73,7 +75,7 @@ pub struct DetectFolds {
 }
 
 /// What the nodes [`DetectFolds`] makes run where the dense algorithm
-/// doesn't apply.
+/// doesn't apply. Where memory runs short, they all run the host's plan.
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 enum Fallback {
     /// The plan each node replaced, so a detected fold is never slower than
@@ -329,7 +331,7 @@ fn fold(agg: &Aggregate, fallback: Fallback) -> Result<Option<LogicalPlan>> {
     };
     let node = match (fallback, masked) {
         (Fallback::HostPlan, _) | (_, true) => node.with_fallback(original),
-        (Fallback::Kernels, false) => node.marked_found(),
+        (Fallback::Kernels, false) => node.with_memory_fallback(original),
     }
     .into_plan();
     // Give the node the aggregate's schema: its names, qualifiers, types
@@ -742,7 +744,7 @@ fn from_domain(p: &Projection, fallback: Fallback) -> Result<Option<LogicalPlan>
     let masked = !fold.conditions().is_empty();
     let new = match (fallback, masked) {
         (Fallback::HostPlan, _) | (_, true) => new.with_fallback(original.clone()),
-        (Fallback::Kernels, false) => new.marked_found(),
+        (Fallback::Kernels, false) => new.with_memory_fallback(original.clone()),
     }
     .into_plan();
     // Rename the node's columns to the join's, so `p` reads them as before.

crates/einfold-datafusion/src/fold_node.rs

@@ -30,6 +30,9 @@ pub struct FoldNode {
     /// The plan this node replaced, if it replaced one: its columns are the
     /// node's, in order, though their names may differ.
     fallback: Option<LogicalPlan>,
+    /// Whether the fallback runs wherever the dense algorithm declines, or
+    /// only where memory runs short.
+    host_on_decline: bool,
     /// Whether [`DetectFolds`](crate::DetectFolds) made this node.
     found: bool,
     /// Whether the node computes its fold only when no two domain rows reach
@@ -78,6 +81,7 @@ impl FoldNode {
             inputs,
             schema,
             fallback: None,
+            host_on_decline: false,
             found: false,
             domain_rows: DomainRows::Any,
             null_when: None,
@@ -87,14 +91,36 @@ impl FoldNode {
     /// This node, keeping `plan` as the plan it replaced. When the dense
     /// algorithm doesn't apply, [`EinFoldExec`](crate::EinFoldExec) runs
     /// `plan` instead of the discrimination algorithm, so the host's own plan is the
-    /// slowest the fold can be.
+    /// slowest the fold can be. It runs `plan` too where the session's
+    /// memory pool can't give it what it needs.
     #[must_use]
     pub fn with_fallback(mut self, plan: LogicalPlan) -> Self {
         self.fallback = Some(plan);
+        self.host_on_decline = true;
         self.found = true;
         self
     }
 
+    /// This node, keeping `plan` as the plan it replaced, which
+    /// [`EinFoldExec`](crate::EinFoldExec) runs only where the session's
+    /// memory pool can't give it what it needs: the host's plan can spill to
+    /// disk, and einfold's kernels can't.
+    #[must_use]
+    pub fn with_memory_fallback(mut self, plan: LogicalPlan) -> Self {
+        self.fallback = Some(plan);
+        self.host_on_decline = false;
+        self.found = true;
+        self
+    }
+
+    /// Whether the node runs its fallback wherever the dense algorithm
+    /// declines (see [`FoldNode::with_fallback`]), and not only where memory
+    /// runs short.
+    #[must_use]
+    pub fn host_on_decline(&self) -> bool {
+        self.host_on_decline
+    }
+
     /// 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,
@@ -127,13 +153,6 @@ impl FoldNode {
         self.null_when
     }
 
-    /// This node, marked as made by [`DetectFolds`](crate::DetectFolds)
-    /// without a fallback.
-    pub(crate) fn marked_found(mut self) -> Self {
-        self.found = true;
-        self
-    }
-
     /// Whether [`DetectFolds`](crate::DetectFolds) made this node.
     pub(crate) fn found(&self) -> bool {
         self.found
@@ -174,6 +193,7 @@ impl PartialOrd for FoldNode {
                 &n.fold,
                 &n.inputs,
                 &n.fallback,
+                n.host_on_decline,
                 n.found,
                 n.domain_rows,
                 n.null_when,
@@ -222,6 +242,7 @@ impl UserDefinedLogicalNodeCore for FoldNode {
         let node = FoldNode::try_new(self.fold.clone(), operands, domain, names.collect())?;
         Ok(FoldNode {
             fallback: self.fallback.clone(),
+            host_on_decline: self.host_on_decline,
             found: self.found,
             domain_rows: self.domain_rows,
             null_when: self.null_when,

crates/einfold-datafusion/src/planner.rs

@@ -52,9 +52,13 @@ impl ExtensionPlanner for EinFoldPlanner {
         };
         let schema = Arc::new(node.schema().as_arrow().clone());
         let mut exec = EinFoldExec::new(node.fold().clone(), physical_inputs.to_vec(), schema);
-        if let Some(plan) = node.fallback() {
-            exec = (exec.with_fallback(plan.clone()))
-                .with_execution_props(session.execution_props().clone());
+        exec = match (node.fallback(), node.host_on_decline()) {
+            (Some(plan), true) => exec.with_fallback(plan.clone()),
+            (Some(plan), false) => exec.with_memory_fallback(plan.clone()),
+            (None, _) => exec,
+        };
+        if node.fallback().is_some() {
+            exec = exec.with_execution_props(session.execution_props().clone());
         }
         if node.domain_rows() == DomainRows::UniqueKeys {
             exec = exec.requiring_unique_domain();

crates/einfold-datafusion/tests/detect_vs_host.rs

@@ -376,6 +376,28 @@ fn einfold(threads: usize, fallback: bool) -> SessionContext {
     SessionContext::new_with_state(state)
 }
 
+/// DataFusion, with einfold or without, under a memory pool of `bytes`
+/// shared fairly among the operators that can spill and given first come,
+/// first served to those that can't, as DataFusion's `FairSpillPool` does.
+fn pooled(threads: usize, with_einfold: Option<DetectFolds>, bytes: usize) -> SessionContext {
+    use datafusion::execution::memory_pool::FairSpillPool;
+    use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+    let config = SessionConfig::new().with_target_partitions(threads);
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::new(FairSpillPool::new(bytes)))
+        .build_arc()
+        .unwrap();
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_config(config)
+        .with_runtime_env(runtime);
+    let state = match with_einfold {
+        Some(rule) => state.with_query_planner(Arc::new(EinFoldQueryPlanner::detecting(rule))),
+        None => state,
+    };
+    SessionContext::new_with_state(state.build())
+}
+
 /// How many `EinFoldExec`s a physical plan has, and how many of them have a
 /// domain.
 fn exec_folds(plan: &Arc<dyn datafusion::physical_plan::ExecutionPlan>) -> (usize, usize) {
@@ -1097,3 +1119,101 @@ async fn masked_folds_over_heads_run_dense() {
     assert!(shown.contains("dense_folds=1"), "{shown}");
     assert!(!shown.contains("host_plans=1"), "{shown}");
 }
+
+/// Property: under a small memory pool, einfold runs out of memory no more
+/// often than the host does. Each case runs a random query of the sweep's
+/// spellings with einfold, by its kernels or with the host's plan as
+/// fallback, and without, each under its own pool of the same size. Where
+/// both answer, they answer the same rows; neither fails but for memory.
+/// einfold reserves what it reads and what its kernels allocate, and where
+/// the pool refuses, runs the host's plan, whose operators can spill. The
+/// operators above einfold's can still differ from the host's, as einfold's
+/// estimates of its output differ from the host's, so a case can run short
+/// of memory in one plan and not the other, either way.
+#[tokio::test]
+async fn einfold_runs_short_of_memory_only_where_the_host_does() {
+    let seed: u64 = std::env::var("EINFOLD_SEED")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(0x5_9001);
+    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 short = |e: &datafusion::error::DataFusionError| {
+        e.to_string().contains("Resources exhausted") || e.to_string().contains("budget")
+    };
+    let (mut both_short, mut host_short, mut ours_short, mut fell_back) = (0, 0, 0, 0);
+    for i in 0..cases {
+        let seed = seed.wrapping_add(i);
+        let mut rng = Rng::new(seed);
+        let keys = [Keys::Dense, Keys::Dense, Keys::Wild, Keys::Text][rng.below(4)];
+        let messy = !matches!(keys, Keys::Dense) || rng.chance(0.4);
+        let rule = match rng.chance(0.5) {
+            true => DetectFolds::with_host_fallback(),
+            false => DetectFolds::default(),
+        };
+        let bytes = [2 << 10, 4 << 10, 8 << 10, 16 << 10, 64 << 10, 256 << 10][rng.below(6)];
+        let threads = 1 + rng.below(4);
+        let (ours, theirs) = (
+            pooled(threads, Some(rule), bytes),
+            pooled(threads, None, bytes),
+        );
+        let ctxs = [&ours, &theirs];
+        let ext = |rng: &mut Rng| 1 + rng.below(6);
+        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);
+        let d_messy = messy || rng.chance(0.3);
+        table(&mut rng, &ctxs, "d", [ea + 1, ec + 1], keys, d_messy);
+        let joined_back = rng.chance(0.2);
+        let (sql, _) = query(&mut rng, joined_back);
+        let fail =
+            |what: String| panic!("seed {seed}, a pool of {bytes} bytes: {what}\nSQL: {sql}");
+        let want = match theirs.sql(&sql).await.unwrap().collect().await {
+            Ok(want) => Some(rows(&want)),
+            Err(e) if short(&e) => None,
+            Err(e) => {
+                fail(format!("the host failed: {e}"));
+                None
+            }
+        };
+        let df = ours.sql(&sql).await.unwrap();
+        let physical = df.create_physical_plan().await.unwrap();
+        let got = datafusion::physical_plan::collect(Arc::clone(&physical), ours.task_ctx()).await;
+        let shown = DisplayableExecutionPlan::with_metrics(physical.as_ref())
+            .indent(true)
+            .to_string();
+        fell_back += usize::from(shown.contains("memory_fallbacks=1"));
+        match (got, want) {
+            (Ok(got), Some(want)) if rows(&got) == want => {}
+            (Ok(got), Some(want)) => {
+                let got = rows(&got);
+                fail(format!("einfold gave\n{got:#?}\nthe host gave\n{want:#?}"));
+            }
+            (Err(e), Some(_)) if short(&e) => ours_short += 1,
+            (Err(e), Some(_)) => {
+                fail(format!("einfold failed where the host answers: {e}"));
+            }
+            (Err(e), None) if short(&e) => both_short += 1,
+            (Err(e), None) => {
+                fail(format!("einfold failed: {e}"));
+            }
+            (Ok(_), None) => host_short += 1,
+        }
+    }
+    eprintln!(
+        "out of memory in both {both_short}, in einfold only {ours_short}, in the host only \
+         {host_short}; einfold ran the host's plan for memory in {fell_back} of {cases} cases"
+    );
+    // Both are claims about many cases: one case, replayed from its seed,
+    // may run short in either.
+    if cases >= 100 {
+        assert!(
+            ours_short <= host_short,
+            "einfold alone ran out of memory in {ours_short} cases, the host alone in {host_short}"
+        );
+        assert!(fell_back > 0, "einfold never ran short of memory");
+    }
+}

crates/einfold-kernels/src/algorithm.rs

@@ -56,17 +56,8 @@ pub fn fold_dense(
     schema: SchemaRef,
     options: Options,
 ) -> Result<Option<Vec<RecordBatch>>> {
-    let threads = options.threads.max(1);
     let domain = domain.map(|d| d.batches);
-    dense::fold(
-        fold,
-        a.batches,
-        b.batches,
-        domain,
-        schema,
-        threads,
-        options.domain_rows,
-    )
+    dense::fold(fold, a.batches, b.batches, domain, schema, &options)
 }
 
 /// `fold` over operands `a` and `b`, and its domain if it has one, as
@@ -92,17 +83,8 @@ pub fn fold(
     if options.domain_rows == DomainRows::UniqueKeys {
         return bad("only fold_dense checks that a domain's keys are unique".into());
     }
-    let threads = options.threads.max(1);
     let domain = domain.map(|d| d.batches);
-    let dense = dense::fold(
-        fold,
-        a.batches,
-        b.batches,
-        domain,
-        schema.clone(),
-        threads,
-        DomainRows::Any,
-    )?;
+    let dense = dense::fold(fold, a.batches, b.batches, domain, schema.clone(), &options)?;
     if let Some(batches) = dense {
         return Ok(Folded {
             algorithm: Algorithm::Dense,
@@ -112,10 +94,16 @@ pub fn fold(
     if !fold.conditions().is_empty() {
         return bad("the discrimination algorithm computes no masks".into());
     }
-    // The discrimination algorithm reads each operand as one batch.
+    // The discrimination algorithm reads each operand as one batch, a copy.
+    let size = |r: &Relation| -> usize {
+        (r.batches.iter())
+            .map(RecordBatch::get_array_memory_size)
+            .sum()
+    };
+    options.reserve(size(&a) + size(&b), "the operands as one batch each")?;
     let whole = |r: Relation| compute::concat_batches(r.schema, r.batches);
     let (a, b) = (whole(a)?, whole(b)?);
-    let batches = discriminate::fold(fold, &a, &b, domain, schema, options.host, threads)?;
+    let batches = discriminate::fold(fold, &a, &b, domain, schema, &options)?;
     Ok(Folded {
         algorithm: Algorithm::Discriminate,
         batches,