From 998d551be794a83efa21c2989431cf72d13b2eec Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 28 Oct 2026 08:35:20 -0700 Subject: [PATCH 1/4] Run the host's plan where memory runs short, as its spilling operators would EinFoldExec read its inputs whole, reserving them from the session's memory pool, and failed where the pool couldn't give them; its kernels' working arrays weren't reserved at all. The host's own operators spill to disk instead, so under a small pool einfold failed queries the host answers about twice as often as the reverse (#74). Now every detected fold keeps the host's plan, which costs nothing until it runs, since it is planned only then. EinFoldExec runs it wherever the pool refuses: while the inputs are read, when the kernels ask for their working arrays, or when the output is reserved. The kernels ask before they allocate, through a Budget in their Options, and stop with a memory error where it says no. The host's plan streams its rows through, as it would in the host's own plan, and einfold's own output gives its reservation back as it is emitted. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/src/detect_folds.rs | 18 +- .../einfold-datafusion/src/ein_fold_exec.rs | 195 ++++++++++++++---- crates/einfold-datafusion/src/fold_node.rs | 37 +++- crates/einfold-datafusion/src/planner.rs | 10 +- .../tests/detect_vs_host.rs | 118 +++++++++++ crates/einfold-kernels/src/algorithm.rs | 32 +-- crates/einfold-kernels/src/dense.rs | 37 +++- crates/einfold-kernels/src/discriminate.rs | 47 ++++- crates/einfold-kernels/src/lib.rs | 2 +- crates/einfold-kernels/src/options.rs | 43 +++- .../einfold-kernels/tests/kernels_vs_sql.rs | 94 ++++++++- 11 files changed, 533 insertions(+), 100 deletions(-) diff --git a/crates/einfold-datafusion/src/detect_folds.rs b/crates/einfold-datafusion/src/detect_folds.rs index 188e2c80..89057a05 100644 --- a/crates/einfold-datafusion/src/detect_folds.rs +++ b/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> { }; 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 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. diff --git a/crates/einfold-datafusion/src/ein_fold_exec.rs b/crates/einfold-datafusion/src/ein_fold_exec.rs index e797842e..34ee7eab 100644 --- a/crates/einfold-datafusion/src/ein_fold_exec.rs +++ b/crates/einfold-datafusion/src/ein_fold_exec.rs @@ -12,6 +12,7 @@ use std::sync::{Arc, OnceLock}; use datafusion::arrow::array::{ArrayRef, BooleanArray, RecordBatch}; use datafusion::arrow::compute::nullif; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::error::ArrowError; use datafusion::common::runtime::SpawnedTask; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{internal_err, DataFusionError, Result}; @@ -21,15 +22,18 @@ use datafusion::logical_expr::execution_props::ExecutionProps; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::metrics::{ + self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, +}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ self, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties, }; use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; use einfold_ir::{ColumnIndex, Fold, Pipeline}; -use einfold_kernels::{Algorithm, DomainRows, Folded, Options, Relation}; -use futures::{future, stream}; +use einfold_kernels::{Algorithm, Budget, DomainRows, Folded, Options, Relation}; +use futures::stream::BoxStream; +use futures::{future, stream, StreamExt}; use crate::{fold_node, pipeline_node}; use crate::{EinFoldPlanner, HOST}; @@ -56,9 +60,13 @@ enum Work { /// fused algorithm. No result depends on the number of threads or on the /// order rows arrive in. /// -/// Like DataFusion's own operators, it reserves the memory its inputs take -/// from the session's memory pool, and fails if the pool can't give it; it -/// runs the kernels on a blocking thread, off the async runtime; and its +/// Like DataFusion's own operators, it reserves the memory it takes from the +/// session's memory pool: its inputs as it reads them, what the kernels +/// allocate before they allocate it, and its output. Where the pool can't +/// give it that, it runs its fallback, if it has one, as the host's spilling +/// operators would have run, and counts it in `memory_fallbacks`; without +/// one, it fails, as an operator that can't spill does. It runs the kernels +/// on a blocking thread, off the async runtime; and its /// metrics, shown by `EXPLAIN ANALYZE`, count which algorithm ran /// (`dense_folds`, `discrimination_folds`, `fused_pipelines`), the pairs of /// tiles a pipeline visited of all it has (`tiles_visited`, `tiles`), its @@ -78,11 +86,14 @@ pub struct EinFoldExec { work: Work, inputs: Vec>, /// The host's plan for the same rows, run instead of the discrimination - /// algorithm, and planned only then. + /// algorithm, or where memory runs short, and planned only then. fallback: Option>, /// The planning session's properties, such as its variable providers, /// for planning the fallback as the session would. execution_props: Option>, + /// Whether the fallback runs wherever the dense algorithm declines, and + /// not only where memory runs short. + host_on_decline: bool, /// Whether to run the fallback unless no two domain rows reach one group. domain_rows: DomainRows, /// The domain's Boolean column that makes a row's value NULL where true. @@ -125,6 +136,7 @@ impl EinFoldExec { inputs, fallback: None, execution_props: None, + host_on_decline: false, domain_rows: DomainRows::Any, null_when: None, properties: Arc::new(properties), @@ -139,6 +151,18 @@ impl EinFoldExec { #[must_use] pub fn with_fallback(mut self, plan: LogicalPlan) -> Self { self.fallback = Some(Arc::new(plan)); + self.host_on_decline = true; + self + } + + /// This operator, running `plan` only where the session's memory pool + /// can't give it what it needs, and the discrimination algorithm where + /// the dense one declines. `plan` must produce the operator's columns, in + /// order, under any names. It is planned only if it runs. + #[must_use] + pub fn with_memory_fallback(mut self, plan: LogicalPlan) -> Self { + self.fallback = Some(Arc::new(plan)); + self.host_on_decline = false; self } @@ -181,6 +205,7 @@ impl EinFoldExec { let mut exec = EinFoldExec::computing(self.work.clone(), inputs, self.schema()); exec.fallback = self.fallback.clone(); exec.execution_props = self.execution_props.clone(); + exec.host_on_decline = self.host_on_decline; exec.domain_rows = self.domain_rows; exec.null_when = self.null_when; exec @@ -289,6 +314,87 @@ async fn run_host_plan( .await } +/// The rows of the host's plan `plan`, planned for the task `ctx` (see +/// [`plan_host`]), as one stream, its columns under this operator's names, +/// `out`, and counted in `output_rows`. They stream through, as they would +/// in the host's own plan, where a Tokio runtime is current; where none is, +/// they are collected on this library's runtime (see [`run_host_plan`]). +async fn host_stream( + plan: Arc, + ctx: Arc, + props: Option>, + out: SchemaRef, + output_rows: metrics::Count, +) -> Result>> { + use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + let rename = move |b: Result| -> Result { + let b = RecordBatch::try_new(Arc::clone(&out), b?.columns().to_vec())?; + output_rows.add(b.num_rows()); + Ok(b) + }; + if tokio::runtime::Handle::try_current().is_err() { + let batches = run_host_plan(plan, ctx, props).await?.concat(); + return Ok(stream::iter(batches.into_iter().map(Ok).map(rename)).boxed()); + } + let physical = plan_host(&plan, &ctx, props.as_deref()).await?; + let one: Arc = match physical.properties().partitioning.partition_count() { + 1 => physical, + _ => Arc::new(CoalescePartitionsExec::new(physical)), + }; + Ok(one.execute(0, ctx)?.map(rename).boxed()) +} + +/// Whether `e` says that memory ran short: that the session's memory pool +/// refused a reservation, or that einfold's kernels found their budget too +/// small. +fn short_of_memory(e: &DataFusionError) -> bool { + match e.find_root() { + DataFusionError::ResourcesExhausted(_) => true, + DataFusionError::ArrowError(e, _) => matches!(**e, ArrowError::MemoryError(_)), + _ => false, + } +} + +/// The session's memory pool, as a budget for einfold's kernels: what they +/// ask for is reserved, as long as `0` is. +#[derive(Debug)] +struct Pool<'a>(&'a MemoryReservation); + +impl Budget for Pool<'_> { + fn grow(&self, bytes: usize) -> bool { + self.0.try_grow(bytes).is_ok() + } +} + +/// `batches` as this operator's output stream: in batches of `batch_size` +/// rows at most, counted in `output_rows`. `reservation` counts them, and +/// gives back each slice's share as it is emitted: from then on, what reads +/// it holds it, and counts it if it keeps it, as it would the host's. +fn slices( + batches: Vec, + batch_size: usize, + output_rows: &metrics::Count, + reservation: Arc, +) -> BoxStream<'static, Result> { + let slices: Vec<(RecordBatch, usize)> = batches + .iter() + .inspect(|b| output_rows.add(b.num_rows())) + .flat_map(|batch| { + let (rows, bytes) = (batch.num_rows(), batch.get_array_memory_size()); + (0..rows).step_by(batch_size).map(move |start| { + let n = batch_size.min(rows - start); + (batch.slice(start, n), bytes * n / rows) + }) + }) + .collect(); + stream::iter(slices) + .map(move |(b, bytes)| { + reservation.shrink(bytes.min(reservation.size())); + Ok(b) + }) + .boxed() +} + /// Every batch of every partition of `plan`, in partition order, each counted /// against `reservation`. The partitions are read concurrently, since an /// operator such as a hash repartition produces them together and would @@ -361,7 +467,7 @@ impl ExecutionPlan for EinFoldExec { partition: usize, context: Arc, ) -> Result { - use futures::{StreamExt, TryStreamExt}; + use futures::TryStreamExt; if partition != 0 { return internal_err!("EinFoldExec has one partition, asked for {partition}"); } @@ -383,6 +489,7 @@ impl ExecutionPlan for EinFoldExec { host: HOST, threads: context.session_config().target_partitions(), domain_rows: self.domain_rows, + budget: None, }; if self.domain_rows == DomainRows::UniqueKeys && self.fallback.is_none() { return internal_err!("EinFoldExec requires a unique domain, but has no fallback"); @@ -394,8 +501,10 @@ impl ExecutionPlan for EinFoldExec { let tiles_visited = builder().global_counter("tiles_visited"); let tiles = builder().global_counter("tiles"); let host_plans = builder().global_counter("host_plans"); + let memory_fallbacks = builder().global_counter("memory_fallbacks"); let fallback = self.fallback.clone(); let execution_props = self.execution_props.clone(); + let host_on_decline = self.host_on_decline; let null_when = self.null_when; let compute = builder().elapsed_compute(0); let output_rows = builder().output_rows(0); @@ -406,12 +515,29 @@ impl ExecutionPlan for EinFoldExec { Arc::new(MemoryConsumer::new("EinFoldExec").register(context.memory_pool())); let out = Arc::clone(&schema); let result = async move { + // The host's plan, where memory ran short: whatever this operator + // held is released first, and the plan's operators reserve their + // own. + let short = |e: DataFusionError| async { + let Some(plan) = fallback.clone().filter(|_| short_of_memory(&e)) else { + return Err(e); + }; + reservation.free(); + memory_fallbacks.add(1); + host_plans.add(1); + let (rows, props) = (output_rows.clone(), execution_props.clone()); + host_stream(plan, Arc::clone(&context), props, Arc::clone(&out), rows).await + }; // Every input is read whole, and stays in the batches it arrived // in: the kernels read them batch by batch. let schemas: Vec = inputs.iter().map(|i| i.schema()).collect(); let reads = (inputs.into_iter()).map(|i| collect(i, Arc::clone(&context), &reservation)); - let mut read = future::try_join_all(reads).await?.into_iter(); + let read = match future::try_join_all(reads).await { + Ok(read) => read, + Err(e) => return short(e).await, + }; + let mut read = read.into_iter(); // Two operands, then a fold's domain or a pipeline's values, as // checked above. let (Some(a), Some(b), third) = (read.next(), read.next(), read.next()) else { @@ -424,11 +550,14 @@ impl ExecutionPlan for EinFoldExec { .zip(domain.as_ref()) .map(|(c, d)| d.iter().map(|b| Arc::clone(b.column(c.0))).collect()); let timer = compute.timer(); - let with_fallback = fallback.is_some(); let kernel_out = Arc::clone(&out); let held = Arc::clone(&reservation); let kernel = move || { - let _held = held; + let pool = Pool(&held); + let options = Options { + budget: Some(&pool), + ..options + }; let (a, b) = ( Relation::new(&schemas[0], &a), Relation::new(&schemas[1], &b), @@ -454,7 +583,7 @@ impl ExecutionPlan for EinFoldExec { }; let domain = domain.as_deref().map(|d| Relation::new(&schemas[2], d)); let tiles = None; - match with_fallback { + match host_on_decline { true => einfold_kernels::fold_dense(fold, a, b, domain, kernel_out, options) .map(|r| { r.map(|batches| { @@ -472,10 +601,14 @@ impl ExecutionPlan for EinFoldExec { let ran = match tokio::runtime::Handle::try_current() { Ok(_) => SpawnedTask::spawn_blocking(kernel) .await - .map_err(|e| DataFusionError::External(Box::new(e)))??, - Err(_) => kernel()?, + .map_err(|e| DataFusionError::External(Box::new(e)))?, + Err(_) => kernel(), }; timer.done(); + let ran = match ran { + Ok(ran) => ran, + Err(e) => return short(DataFusionError::from(e)).await, + }; // The kernel has dropped its inputs: release them, so that the // host's plan, if it runs, and the operators downstream have the // memory they would have had without einfold. The output, which @@ -498,7 +631,7 @@ impl ExecutionPlan for EinFoldExec { }) .collect() }; - let batches = match (ran, fallback) { + let batches = match (ran, fallback.clone()) { (Some((Folded { algorithm, batches }, counts)), _) => { if let Some((visited, all)) = counts { tiles_visited.add(visited); @@ -514,34 +647,20 @@ impl ExecutionPlan for EinFoldExec { masked(batches)? } (None, Some(plan)) => { - // The host's plan, its columns under this operator's names. host_plans.add(1); - run_host_plan(plan, Arc::clone(&context), execution_props.clone()) - .await? - .concat() - .into_iter() - .map(|b| RecordBatch::try_new(Arc::clone(&out), b.columns().to_vec())) - .collect::, _>>()? + let rows = output_rows.clone(); + let props = execution_props.clone(); + return host_stream(plan, Arc::clone(&context), props, Arc::clone(&out), rows) + .await; } (None, None) => return internal_err!("EinFoldExec ran no algorithm"), }; - for b in &batches { - output_rows.add(b.num_rows()); - reservation.try_grow(b.get_array_memory_size())?; + let size: usize = batches.iter().map(RecordBatch::get_array_memory_size).sum(); + if let Err(e) = reservation.try_grow(size) { + drop(batches); + return short(e).await; } - let slices: Vec> = batches - .iter() - .flat_map(|batch| { - (0..batch.num_rows()).step_by(batch_size).map(move |start| { - Ok(batch.slice(start, batch_size.min(batch.num_rows() - start))) - }) - }) - .collect(); - // The reservation lives as long as the output does. - Ok::<_, DataFusionError>(stream::iter(slices).map(move |b| { - let _ = &reservation; - b - })) + Ok::<_, DataFusionError>(slices(batches, batch_size, &output_rows, reservation)) }; let stream = stream::once(result).try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new( diff --git a/crates/einfold-datafusion/src/fold_node.rs b/crates/einfold-datafusion/src/fold_node.rs index 464a494e..1bfe4df3 100644 --- a/crates/einfold-datafusion/src/fold_node.rs +++ b/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, + /// 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, diff --git a/crates/einfold-datafusion/src/planner.rs b/crates/einfold-datafusion/src/planner.rs index 1ef78894..cab56c6b 100644 --- a/crates/einfold-datafusion/src/planner.rs +++ b/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(); diff --git a/crates/einfold-datafusion/tests/detect_vs_host.rs b/crates/einfold-datafusion/tests/detect_vs_host.rs index bd9646ec..eb5a56e1 100644 --- a/crates/einfold-datafusion/tests/detect_vs_host.rs +++ b/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, 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) -> (usize, usize) { @@ -1097,3 +1119,99 @@ 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" + ); + assert!( + ours_short <= host_short, + "einfold alone ran out of memory in {ours_short} cases, the host alone in {host_short}" + ); + if cases >= 100 { + assert!(fell_back > 0, "einfold never ran short of memory"); + } +} diff --git a/crates/einfold-kernels/src/algorithm.rs b/crates/einfold-kernels/src/algorithm.rs index 6681f5e7..1ef851c8 100644 --- a/crates/einfold-kernels/src/algorithm.rs +++ b/crates/einfold-kernels/src/algorithm.rs @@ -56,17 +56,8 @@ pub fn fold_dense( schema: SchemaRef, options: Options, ) -> Result>> { - 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, diff --git a/crates/einfold-kernels/src/dense.rs b/crates/einfold-kernels/src/dense.rs index 6c74eebb..771ba027 100644 --- a/crates/einfold-kernels/src/dense.rs +++ b/crates/einfold-kernels/src/dense.rs @@ -62,7 +62,7 @@ use arrow::compute; use arrow::datatypes::{DataType, SchemaRef}; use einfold_ir::{ColumnIndex, DimensionId, Domain, Fold, Operand}; -use crate::options::DomainRows; +use crate::options::{DomainRows, Options}; use crate::Result; /// Output rows per block, when the output has at least [`FEW_CELLS`] cells. @@ -179,16 +179,20 @@ pub(crate) fn fold( b: &[RecordBatch], domain: Option<&[RecordBatch]>, schema: SchemaRef, - threads: usize, - domain_rows: DomainRows, + options: &Options, ) -> Result>> { let [op_a, op_b] = fold.operands() else { return Ok(None); }; + let domain_rows = options.domain_rows; let inputs: usize = (a.iter().chain(b).chain(domain.unwrap_or_default())) .map(RecordBatch::num_rows) .sum(); - let threads = if inputs < FEW_ROWS { 1 } else { threads }; + let threads = if inputs < FEW_ROWS { + 1 + } else { + options.threads.max(1) + }; // The roles: W, X and Y in output order, Z in dimension order. Here a // dimension is its index, which indexes the extents too. let (mut ws, mut xs, mut ys, mut zs) = (Vec::new(), Vec::new(), Vec::new(), Vec::new()); @@ -225,6 +229,10 @@ pub(crate) fn fold( // Each batch's key columns as non-negative i64, cast once, with their // extents, batches in parallel; then each dimension's extent, its largest // key in either operand, plus one. + let rows = |batches: &[RecordBatch]| batches.iter().map(|b| b.num_rows()).sum::(); + // The keys, cast: 8 bytes a key. + let cast = rows(a) * op_a.keys.len() + rows(b) * op_b.keys.len(); + options.reserve(cast.saturating_mul(8), "the dense algorithm's keys")?; let cast_keys = |op: &Operand, batches: &[RecordBatch]| { let batch_keys = par_tasks(batches.len(), threads, |i| { (op.keys.iter()) @@ -251,7 +259,6 @@ pub(crate) fn fold( .collect() }; let (a_keys, b_keys) = (strip(a_keys), strip(b_keys)); - let rows = |batches: &[RecordBatch]| batches.iter().map(|b| b.num_rows()).sum::(); let size = |dims: &[usize]| { dims.iter() .try_fold(1usize, |n, &d| n.checked_mul(extent[d])) @@ -273,6 +280,26 @@ pub(crate) fn fold( { return Ok(None); } + // What the algorithm allocates from here on: a's matrix, at 9 bytes a + // cell, and its rows' flags; b's matrices, likewise; each row's two + // positions, at 4 bytes each; the output as it is accumulated, and its + // copy; and the output's columns, 8 bytes each, with each domain row's + // cell. + let (na, nb) = ( + nx.saturating_mul(nz), + nw.saturating_mul(nz).saturating_mul(ny), + ); + let out_columns = fold.outputs().len() + 1 + usize::from(domain.is_some()); + let working = [ + na.saturating_mul(9), + nx.saturating_mul(2), + nb.saturating_mul(9), + (rows(a) + rows(b)).saturating_mul(8), + nx.saturating_mul(ny).saturating_mul(16), + out_rows.saturating_mul(out_columns).saturating_mul(8), + ]; + let working = working.into_iter().fold(0usize, usize::saturating_add); + options.reserve(working, "the dense algorithm's matrices")?; let factors = |op: &Operand, batch: &'_ RecordBatch| { op.factor .and_then(|c| batch.column(c.0).as_any().downcast_ref::()) diff --git a/crates/einfold-kernels/src/discriminate.rs b/crates/einfold-kernels/src/discriminate.rs index 4814ca50..2b18e09b 100644 --- a/crates/einfold-kernels/src/discriminate.rs +++ b/crates/einfold-kernels/src/discriminate.rs @@ -51,7 +51,8 @@ use arrow::error::ArrowError; use einfold_ir::{DimensionId, Fold, KeyEquality}; use crate::keys::{self, Keys}; -use crate::{dense, Host, Result}; +use crate::options::Options; +use crate::{dense, Result}; /// Output keys of `a` per block of the product, each block on one thread. const XB: usize = 256; @@ -371,9 +372,9 @@ pub(crate) fn fold( b: &RecordBatch, domain: Option<&[RecordBatch]>, schema: SchemaRef, - host: Host, - threads: usize, + options: &Options, ) -> Result> { + let (host, threads) = (options.host, options.threads.max(1)); // Rows are ranked as u32s: the operands' together, and each operand's // with the domain's. let (na, nb) = (a.num_rows(), b.num_rows()); @@ -410,6 +411,30 @@ pub(crate) fn fold( f.ok_or_else(|| ArrowError::InvalidArgumentError("factors must be Float64".into())) }; let (a_factors, b_factors) = (factors(0, a)?, factors(1, b)?); + // What ranking and ordering the rows allocate, per row ranked: its keys + // in the row format, which take about what the key columns do; the + // radix sort's words and pairs, 24 bytes; its rank, 4; and the orders of + // each operand's rows, with b's laid out by z, about 40 more. The join's + // keys rank a's and b's rows, and the outputs' rank them again, with the + // domain's. + let nd: usize = domain.map_or(0, |d| d.iter().map(RecordBatch::num_rows).sum()); + let key_bytes = |batch: &RecordBatch, op: usize| -> usize { + (ops[op].keys.iter()) + .map(|k| batch.column(k.column.0).get_array_memory_size()) + .sum() + }; + let ranked = (na + nb) + .saturating_mul(2) + .saturating_add(nd.saturating_mul(2)); + let working = ranked + .saturating_mul(28) + .saturating_add((na + nb).saturating_mul(40)) + .saturating_add( + key_bytes(a, 0) + .saturating_add(key_bytes(b, 1)) + .saturating_mul(2), + ); + options.reserve(working, "the discrimination algorithm's ranks")?; // The join's keys, over both operands at once; rows with a NULL in a key // compared with `=` join nothing. @@ -502,6 +527,22 @@ pub(crate) fn fold( }) .collect(); + // The groups the product makes: no more than the joined pairs, nor than + // every (x, y). Each takes 24 bytes, and then, without a domain, a row of + // the output, 8 bytes a column; with one, the output is a column of 8 + // bytes per domain row, beside the domain's own columns. + let pairs: usize = (a_rows.iter()) + .map(|&r| b_start[a_z(r) + 1] - b_start[a_z(r)]) + .fold(0, usize::saturating_add); + let groups = pairs.min(nx.saturating_mul(ny)); + let output = match fold.domain() { + None => groups.saturating_mul(8 * (fold.outputs().len() + 1)), + Some(_) => nd.saturating_mul(8), + }; + options.reserve( + groups.saturating_mul(24).saturating_add(output), + "the discrimination algorithm's groups", + )?; // Gustavson's product, blocks of x on the threads in turn, each block's // groups in (x, y) order. let blocks = dense::par_tasks(nx.div_ceil(XB), threads, |t| { diff --git a/crates/einfold-kernels/src/lib.rs b/crates/einfold-kernels/src/lib.rs index 8a78c393..7ce89f9b 100644 --- a/crates/einfold-kernels/src/lib.rs +++ b/crates/einfold-kernels/src/lib.rs @@ -39,7 +39,7 @@ mod schema; pub use algorithm::{fold, fold_dense, pipeline, Algorithm, Folded}; pub use fused::Piped; -pub use options::{DomainRows, FloatKeys, Host, Options}; +pub use options::{Budget, DomainRows, FloatKeys, Host, Options}; pub use relation::Relation; pub use schema::{pipeline_schema, schema}; diff --git a/crates/einfold-kernels/src/options.rs b/crates/einfold-kernels/src/options.rs index 61fe632c..a6e56494 100644 --- a/crates/einfold-kernels/src/options.rs +++ b/crates/einfold-kernels/src/options.rs @@ -4,6 +4,10 @@ //! What a caller tells the kernels about the host, and how to run a fold. +use std::fmt; + +use arrow::error::ArrowError; + /// How the host compares floating-point keys, in joins and in `GROUP BY`. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord)] #[non_exhaustive] @@ -39,24 +43,55 @@ pub enum DomainRows { UniqueKeys, } +/// The memory a kernel may take beyond its inputs, such as a query engine's +/// memory pool. A kernel asks before it allocates its working arrays and its +/// output, and stops with [`ArrowError::MemoryError`] where the budget says +/// no, before allocating, so that its caller can run something that needs +/// less, such as a plan that spills to disk. +/// +/// What a kernel asks for is an estimate of what it allocates, within a +/// small factor; it is not given back as the kernel frees it. +/// +/// [`ArrowError::MemoryError`]: arrow::error::ArrowError::MemoryError +pub trait Budget: Sync + fmt::Debug { + /// Whether `bytes` more may be allocated; if so, they are counted. + fn grow(&self, bytes: usize) -> bool; +} + /// How to run a fold. The default is one thread, a host that compares float -/// keys by their bits, and any domain rows. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Options { +/// keys by their bits, any domain rows, and no limit on memory. +#[derive(Clone, Copy, Debug)] +pub struct Options<'a> { /// What the host decides. pub host: Host, /// The most threads to use. The result's bits don't depend on it. pub threads: usize, /// Whether the domain's rows must have unique keys. pub domain_rows: DomainRows, + /// The memory the kernel may take beyond its inputs, or `None` for no + /// limit. + pub budget: Option<&'a dyn Budget>, +} + +impl Options<'_> { + /// Asks the budget for `bytes` to hold `what`, failing if it says no. + pub(crate) fn reserve(&self, bytes: usize, what: &str) -> crate::Result<()> { + match self.budget { + Some(budget) if !budget.grow(bytes) => Err(ArrowError::MemoryError(format!( + "EinFold's budget can't give {bytes} bytes for {what}" + ))), + _ => Ok(()), + } + } } -impl Default for Options { +impl Default for Options<'_> { fn default() -> Self { Options { host: Host::default(), threads: 1, domain_rows: DomainRows::Any, + budget: None, } } } diff --git a/crates/einfold-kernels/tests/kernels_vs_sql.rs b/crates/einfold-kernels/tests/kernels_vs_sql.rs index 534a8a8f..483ac008 100644 --- a/crates/einfold-kernels/tests/kernels_vs_sql.rs +++ b/crates/einfold-kernels/tests/kernels_vs_sql.rs @@ -27,6 +27,10 @@ //! 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. +//! 5. **Budgets.** Given a memory budget of exactly the bytes the kernels +//! asked for with none, they give the same rows by the same algorithm, bit +//! for bit; given one byte fewer, they stop with a memory error. So a +//! kernel asks before it allocates, and asks the same for the same rows. //! //! 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 @@ -40,6 +44,7 @@ //! `EINFOLD_CASES=` (and a starting `EINFOLD_SEED`). use std::collections::HashSet; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use datafusion::arrow::array::{ @@ -47,6 +52,7 @@ use datafusion::arrow::array::{ StringArray, UInt64Array, }; use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::error::ArrowError; use datafusion::common::ScalarValue; use datafusion::datasource::MemTable; use datafusion::prelude::{SessionConfig, SessionContext}; @@ -54,7 +60,7 @@ use einfold_ir::{ Aggregate, ColumnIndex, Dimension, DimensionId, Domain, DomainKey, Fold, Key, KeyEquality, Operand, Output, RowValue, }; -use einfold_kernels::{Algorithm, DomainRows, FloatKeys, Host, Options, Relation}; +use einfold_kernels::{Algorithm, Budget, DomainRows, FloatKeys, Folded, Host, Options, Relation}; struct Rng(u64); @@ -657,9 +663,80 @@ fn run_fold(case: &Case, rng: &mut Rng) -> (Vec, Algorithm) { }; let (a, b) = (batches(&case.a, rng), batches(&case.b, rng)); let domain = case.domain.as_ref().map(|t| batches(t, rng)); + let threads = 1 + rng.below(8); + let folded = fold_within(case, [&a, &b], domain.as_deref(), threads, None).unwrap(); + (folded.batches, folded.algorithm) +} + +/// A budget that counts what it gives, up to a limit. +#[derive(Debug)] +struct Limit { + given: AtomicUsize, + limit: usize, +} + +impl Budget for Limit { + fn grow(&self, bytes: usize) -> bool { + let given = self.given.load(Ordering::Relaxed); + let ok = given.saturating_add(bytes) <= self.limit; + if ok { + self.given.fetch_add(bytes, Ordering::Relaxed); + } + ok + } +} + +/// Checks invariant 5 for a case, on one batching of its tables: the bytes +/// the kernels ask for, and the result given exactly those and one fewer. +fn check_budget(case: &Case, rng: &mut Rng, fail: impl Fn(String)) { + let batches = |t: &Table, rng: &mut Rng| -> Vec { + t.partitions(rng).into_iter().flatten().collect() + }; + let (a, b) = (batches(&case.a, rng), batches(&case.b, rng)); + let domain = case.domain.as_ref().map(|t| batches(t, rng)); + let threads = 1 + rng.below(8); + let run = |limit: usize| { + let budget = Limit { + given: AtomicUsize::new(0), + limit, + }; + let folded = fold_within(case, [&a, &b], domain.as_deref(), threads, Some(&budget)); + (folded, budget.given.into_inner()) + }; + let (free, asked) = run(usize::MAX); + let free = free.unwrap(); + let (exact, _) = run(asked); + match exact { + Ok(f) if f.algorithm == free.algorithm && same_bits(&f.batches, &free.batches) => {} + other => fail(format!("given the {asked} bytes it asked for: {other:?}")), + } + if asked > 0 { + match run(asked - 1).0 { + Err(ArrowError::MemoryError(_)) => {} + other => fail(format!( + "given {} of the {asked} bytes it asked for: {other:?}", + asked - 1 + )), + } + } +} +/// Whether two results have the same rows, in the same order, bit for bit. +fn same_bits(x: &[RecordBatch], y: &[RecordBatch]) -> bool { + x.len() == y.len() && x.iter().zip(y).all(|(p, q)| p == q) +} + +/// `case`'s fold over operand batches `a` and `b`, and `domain`'s, on +/// `threads` threads, within `budget`. +fn fold_within( + case: &Case, + [a, b]: [&[RecordBatch]; 2], + domain: Option<&[RecordBatch]>, + threads: usize, + budget: Option<&dyn Budget>, +) -> Result { let domain = (case.domain.as_ref()) - .zip(domain.as_deref()) + .zip(domain) .map(|(t, d)| Relation::new(&t.schema, d)); let schema = einfold_kernels::schema( &case.fold, @@ -669,19 +746,18 @@ fn run_fold(case: &Case, rng: &mut Rng) -> (Vec, Algorithm) { .unwrap(); let options = Options { host: HOST, - threads: 1 + rng.below(8), + threads, + budget, ..Options::default() }; - let folded = einfold_kernels::fold( + einfold_kernels::fold( &case.fold, - Relation::new(&case.a.schema, &a), - Relation::new(&case.b.schema, &b), + Relation::new(&case.a.schema, a), + Relation::new(&case.b.schema, b), domain, schema, options, ) - .unwrap(); - (folded.batches, folded.algorithm) } /// Checks invariant 4 for a case with a domain that the dense path computes, @@ -705,6 +781,7 @@ fn check_unique_domain(case: &Case, rng: &mut Rng, fail: impl Fn(String)) -> boo host: HOST, threads: 1 + rng.below(8), domain_rows, + budget: None, }; let d = d.map(|batches| Relation::new(&domain.schema, batches)); einfold_kernels::fold_dense(fold, a, b, d, schema.unwrap(), options).unwrap() @@ -789,6 +866,7 @@ async fn fold_matches_sql() { let (ours, ran) = run_fold(&case, &mut rng); let (again, _) = run_fold(&case, &mut rng); + check_budget(&case, &mut rng, fail); dense_family += usize::from(dense); dense_ran += usize::from(ran == Algorithm::Dense); let [op_a, op_b] = case.fold.operands() else { From 04c201c8bf3c12e3737804617e056022e0b05cff Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Wed, 28 Oct 2026 23:11:46 -0700 Subject: [PATCH 2/4] Give the dense fuzz test's options no budget Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-kernels/tests/dense_metamorphic.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/einfold-kernels/tests/dense_metamorphic.rs b/crates/einfold-kernels/tests/dense_metamorphic.rs index 2dda0988..f50ec423 100644 --- a/crates/einfold-kernels/tests/dense_metamorphic.rs +++ b/crates/einfold-kernels/tests/dense_metamorphic.rs @@ -517,6 +517,7 @@ fn run( host, threads, domain_rows, + budget: None, }; let f = einfold_kernels::fold_dense(&fold, a, b, domain, schema, options).expect("runs"); From da6ebd0d7a796fdeaf3cbe21034458719c202c37 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Thu, 29 Oct 2026 12:30:50 -0700 Subject: [PATCH 3/4] Release a batch's reservation with its last slice, and decline other keys before reserving Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/src/ein_fold_exec.rs | 9 ++++++--- crates/einfold-kernels/src/dense.rs | 10 ++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/crates/einfold-datafusion/src/ein_fold_exec.rs b/crates/einfold-datafusion/src/ein_fold_exec.rs index 34ee7eab..a815e3d7 100644 --- a/crates/einfold-datafusion/src/ein_fold_exec.rs +++ b/crates/einfold-datafusion/src/ein_fold_exec.rs @@ -368,8 +368,10 @@ impl Budget for Pool<'_> { /// `batches` as this operator's output stream: in batches of `batch_size` /// rows at most, counted in `output_rows`. `reservation` counts them, and -/// gives back each slice's share as it is emitted: from then on, what reads -/// it holds it, and counts it if it keeps it, as it would the host's. +/// gives back each batch's share as its last slice is emitted: the slices +/// share the batch's buffers, which this stream holds until then. From then +/// on, what reads it holds it, and counts it if it keeps it, as it would the +/// host's. fn slices( batches: Vec, batch_size: usize, @@ -383,7 +385,8 @@ fn slices( let (rows, bytes) = (batch.num_rows(), batch.get_array_memory_size()); (0..rows).step_by(batch_size).map(move |start| { let n = batch_size.min(rows - start); - (batch.slice(start, n), bytes * n / rows) + let last = start + n == rows; + (batch.slice(start, n), if last { bytes } else { 0 }) }) }) .collect(); diff --git a/crates/einfold-kernels/src/dense.rs b/crates/einfold-kernels/src/dense.rs index 771ba027..31a7021f 100644 --- a/crates/einfold-kernels/src/dense.rs +++ b/crates/einfold-kernels/src/dense.rs @@ -230,6 +230,16 @@ pub(crate) fn fold( // extents, batches in parallel; then each dimension's extent, its largest // key in either operand, plus one. let rows = |batches: &[RecordBatch]| batches.iter().map(|b| b.num_rows()).sum::(); + // Keys that aren't integers are declined below; declined before asking + // for memory, so that they don't spend the budget. + let integers = |op: &Operand, batches: &[RecordBatch]| { + (batches.first()).is_none_or(|batch| { + (op.keys.iter()).all(|key| batch.column(key.column.0).data_type().is_integer()) + }) + }; + if !integers(op_a, a) || !integers(op_b, b) { + return Ok(None); + } // The keys, cast: 8 bytes a key. let cast = rows(a) * op_a.keys.len() + rows(b) * op_b.keys.len(); options.reserve(cast.saturating_mul(8), "the dense algorithm's keys")?; From 3dd56ae157360b6a29e7049d1fb824b1491cc2ba Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Thu, 29 Oct 2026 14:01:59 -0700 Subject: [PATCH 4/4] Assert the memory comparison only over many cases Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M96xSez6PZwTmfaGh2eJgS --- crates/einfold-datafusion/tests/detect_vs_host.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/crates/einfold-datafusion/tests/detect_vs_host.rs b/crates/einfold-datafusion/tests/detect_vs_host.rs index eb5a56e1..849f5833 100644 --- a/crates/einfold-datafusion/tests/detect_vs_host.rs +++ b/crates/einfold-datafusion/tests/detect_vs_host.rs @@ -1207,11 +1207,13 @@ async fn einfold_runs_short_of_memory_only_where_the_host_does() { "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" ); - assert!( - ours_short <= host_short, - "einfold alone ran out of memory in {ours_short} cases, the host alone in {host_short}" - ); + // 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"); } }