Carry an extent hint on batch returns

This commit is contained in:
Dennis Kobert
2026-08-21 12:32:57 +00:00
parent 61cf619f4b
commit e68254eff3
6 changed files with 46 additions and 20 deletions

View File

@@ -7,11 +7,15 @@ use std::ops::Range;
#[derive(Debug)]
pub enum BatchStatus<'a> {
/// Producer-resident lanes, shared: read-only for the caller.
Lent(RecordBatch<'a>, Finality),
/// Producer-resident lanes, shared: read-only for the caller. The extent
/// is the producer's knowledge of the level's total after serving the
/// range: a sound lower bound, or exact; a batch shorter than the
/// requested range carries `Exactly` and marks the end of the data.
Lent(RecordBatch<'a>, Finality, Extent),
/// The caller's scratch, filled: the caller is the exclusive owner and may
/// mutate the lanes or reclaim the buffer for in-place reuse.
Filled(RecordBatchMut<'a>, Finality),
/// mutate the lanes or reclaim the buffer for in-place reuse. The extent
/// hint is as for `Lent`.
Filled(RecordBatchMut<'a>, Finality, Extent),
/// No batch implementation behind this edge; a driver answers with the
/// per-lane eval and copy-out loop ([`crate::record::fill_frames`]).
Unbatched,

View File

@@ -499,6 +499,8 @@ where
let base = scratch.as_mut_ptr().cast::<u8>();
let mut local = *input;
let mut finality = crate::gpoll::Finality::AllFinal;
let mut filled = len;
let mut hint = crate::gpoll::Extent::AtLeast(range.end as usize);
for lane in 0..len {
local.set_index(range.start + lane as u64);
let mark = stack::sp();
@@ -510,6 +512,16 @@ where
}
GPoll::Pending => return BatchStatus::Pending,
GPoll::Fallback(boxed) => return BatchStatus::Error(boxed.1),
// A lane past a lower-bound level ends the data: the fill comes
// back short and the hint turns exact.
GPoll::Error(error) if error.kind == crate::gpoll::ErrorKind::PastEnd => {
filled = lane;
hint = crate::gpoll::Extent::Exactly(range.start as usize + lane);
// SAFETY: the failed lane produced no record, so nothing above
// its mark is live.
unsafe { stack::rewind(mark) };
break;
}
GPoll::Error(error) => return BatchStatus::Error(*error),
};
// SAFETY: the lane region is in-bounds by the scratch check, and the
@@ -519,8 +531,8 @@ where
stack::rewind(mark);
}
}
// SAFETY: all `len` lanes were filled above with records of `layout`.
BatchStatus::Filled(unsafe { crate::node::RecordBatchMut::new(scratch, len, layout) }, finality)
// SAFETY: the first `filled` lanes were filled above with records of `layout`.
BatchStatus::Filled(unsafe { crate::node::RecordBatchMut::new(scratch, filled, layout) }, finality, hint)
}
/// The driver a consumer runs on a record edge: a resident batch returns with

View File

@@ -1145,8 +1145,8 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
};
let __start: u64 = #start;
let __batch = match #core_types::record::materialize_batch(&self.#name, __input, __start..__start + __count as u64, __arena) {
#core_types::node::BatchStatus::Lent(__batch, _) => __batch,
#core_types::node::BatchStatus::Filled(__batch, _) => __batch.into_shared(),
#core_types::node::BatchStatus::Lent(__batch, ..) => __batch,
#core_types::node::BatchStatus::Filled(__batch, ..) => __batch.into_shared(),
#core_types::node::BatchStatus::Pending => #pending,
#core_types::node::BatchStatus::Error(__error) => #batch_error,
_ => #batch_failed,
@@ -1343,8 +1343,8 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
_ => return #core_types::gpoll::GPoll::Error(::std::boxed::Box::new(#core_types::gpoll::GraphError::new("extent over a non-exact ranked input"))),
};
match #core_types::record::materialize_batch(&self.#name, __input, 0..__count as u64, __arena) {
#core_types::node::BatchStatus::Lent(__batch, _) => #core_types::gpoll::GPoll::Final(unsafe { #core_types::node::List::<#ty>::new(__batch) }),
#core_types::node::BatchStatus::Filled(__batch, _) => #core_types::gpoll::GPoll::Final(unsafe { #core_types::node::List::<#ty>::new(__batch.into_shared()) }),
#core_types::node::BatchStatus::Lent(__batch, ..) => #core_types::gpoll::GPoll::Final(unsafe { #core_types::node::List::<#ty>::new(__batch) }),
#core_types::node::BatchStatus::Filled(__batch, ..) => #core_types::gpoll::GPoll::Final(unsafe { #core_types::node::List::<#ty>::new(__batch.into_shared()) }),
#core_types::node::BatchStatus::Pending => #core_types::gpoll::GPoll::Pending,
#core_types::node::BatchStatus::Error(__error) => #core_types::gpoll::GPoll::Error(::std::boxed::Box::new(__error)),
_ => #core_types::gpoll::GPoll::Error(::std::boxed::Box::new(#core_types::gpoll::GraphError::new("extent could not materialize a ranked input"))),
@@ -1938,6 +1938,8 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
#(#hoisted_clamps)*
let __frames = __scratch.as_mut_ptr().cast::<u8>();
let mut __finality = #core_types::gpoll::Finality::AllFinal;
let mut __filled = __len;
let mut __hint = #core_types::gpoll::Extent::AtLeast(__range.end as usize);
let mut __lane_ctx = __base_ctx;
for __lane in 0..__len {
#core_types::context::InjectIndex::set_index(&mut __lane_ctx, __range.start + __lane as u64);
@@ -1955,6 +1957,14 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
}
#core_types::gpoll::GPoll::Pending => return #core_types::node::BatchStatus::Pending,
#core_types::gpoll::GPoll::Fallback(__boxed) => return #core_types::node::BatchStatus::Error(__boxed.1),
// A lane past a lower-bound level ends the data: the fill
// comes back short and the hint turns exact.
#core_types::gpoll::GPoll::Error(__error) if __error.kind == #core_types::gpoll::ErrorKind::PastEnd => {
__filled = __lane;
__hint = #core_types::gpoll::Extent::Exactly(__range.start as usize + __lane);
unsafe { #core_types::record::stack::rewind(__lane_mark) };
break;
}
#core_types::gpoll::GPoll::Error(__error) => return #core_types::node::BatchStatus::Error(*__error),
};
// SAFETY: the lane region is in-bounds by the scratch check,
@@ -1967,9 +1977,9 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
// SAFETY: every lane was copied into the caller's scratch, so
// nothing above the entry mark is live.
unsafe { #core_types::record::stack::rewind(__entry_mark) };
// SAFETY: all `__len` lanes were filled above with records of
// the node's layout.
#core_types::node::BatchStatus::Filled(unsafe { #core_types::node::RecordBatchMut::new(__scratch, __len, __node_layout) }, __finality)
// SAFETY: the first `__filled` lanes were filled above with
// records of the node's layout.
#core_types::node::BatchStatus::Filled(unsafe { #core_types::node::RecordBatchMut::new(__scratch, __filled, __node_layout) }, __finality, __hint)
}
}
}

View File

@@ -1330,7 +1330,7 @@ mod tests {
assert!(matches!(node.eval_batch(&scoped, 0..6, None), core_types::node::BatchStatus::NeedBuffer));
let mut scratch = vec![std::mem::MaybeUninit::<u64>::uninit(); 6 * out.lane_stride() / 8];
let core_types::node::BatchStatus::Filled(batch, finality) = node.eval_batch(&scoped, 0..6, Some(&mut scratch)) else {
let core_types::node::BatchStatus::Filled(batch, finality, _) = node.eval_batch(&scoped, 0..6, Some(&mut scratch)) else {
panic!("expected a filled batch");
};
assert_eq!(finality, core_types::gpoll::Finality::AllFinal);
@@ -1389,7 +1389,7 @@ mod tests {
let scoped = ctx.promoted(&head, 0);
let mut scratch = vec![std::mem::MaybeUninit::<u64>::uninit(); 6 * out.lane_stride() / 8];
let core_types::node::BatchStatus::Filled(batch, _) = node.eval_batch(&scoped, 0..6, Some(&mut scratch)) else {
let core_types::node::BatchStatus::Filled(batch, ..) = node.eval_batch(&scoped, 0..6, Some(&mut scratch)) else {
panic!("expected a filled batch");
};
assert_eq!(batch.len(), 6);

View File

@@ -171,8 +171,8 @@ where
_ => return Err(GPoll::error("map content extent is not exact")),
};
match materialize_batch(content, ctx, 0..count as u64, arena) {
BatchStatus::Lent(batch, _) => Ok(unsafe { core_types::node::List::new(batch) }),
BatchStatus::Filled(batch, _) => Ok(unsafe { core_types::node::List::new(batch.into_shared()) }),
BatchStatus::Lent(batch, ..) => Ok(unsafe { core_types::node::List::new(batch) }),
BatchStatus::Filled(batch, ..) => Ok(unsafe { core_types::node::List::new(batch.into_shared()) }),
BatchStatus::Pending => Err(GPoll::Pending),
BatchStatus::Error(error) => Err(GPoll::Error(Box::new(error))),
_ => Err(GPoll::error("map content could not materialize")),
@@ -533,7 +533,7 @@ mod tests {
let scoped = ctx.promoted(&head, 0);
let mut scratch = vec![std::mem::MaybeUninit::<u64>::uninit(); 5 * out.lane_stride() / 8];
let core_types::node::BatchStatus::Filled(batch, _) = node.eval_batch(&scoped, 0..5, Some(&mut scratch)) else {
let core_types::node::BatchStatus::Filled(batch, ..) = node.eval_batch(&scoped, 0..5, Some(&mut scratch)) else {
panic!("expected a filled batch");
};
let batch = batch.into_shared();
@@ -643,7 +643,7 @@ mod tests {
let scoped = ctx.promoted(&head, 0);
let mut scratch = vec![std::mem::MaybeUninit::<u64>::uninit(); 3 * out.lane_stride() / 8];
let core_types::node::BatchStatus::Filled(batch, _) = node.eval_batch(&scoped, 0..3, Some(&mut scratch)) else {
let core_types::node::BatchStatus::Filled(batch, ..) = node.eval_batch(&scoped, 0..3, Some(&mut scratch)) else {
panic!("expected a filled batch");
};
let batch = batch.into_shared();

View File

@@ -1143,7 +1143,7 @@ mod graphene_test {
// One u64 word per lane: the uninstalled layout keeps the f64 inline.
let mut scratch = [const { MaybeUninit::uninit() }; 4];
let status = erased.eval_batch(&ctx, 2..6, Some(&mut scratch));
let BatchStatus::Filled(batch, finality) = status else {
let BatchStatus::Filled(batch, finality, _) = status else {
panic!("expected filled, got {status:?}");
};
let mut got = Vec::new();