fix(promql): charge binary buffers at actual allocation boundaries

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-21 13:12:21 +08:00
parent 90f08fd0a1
commit 4b7f7fc3ab
3 changed files with 309 additions and 34 deletions
+42
View File
@@ -5051,3 +5051,45 @@ fn a_dense_wide_result_fits_a_32_mib_pool() {
assert_eq!(pool.reserved(), 0);
assert_eq!(execution.state().reserved_bytes(), 0);
}
/// An all-missing match group of two right series finishes in a 4 MiB pool and answers nothing.
///
/// No member carries a sample, so the duplicate scan needs no point buffer.
#[test]
fn an_all_missing_match_group_fits_its_pool() {
let points = 100_000;
let (tree, grid, left, right) = binary_tree(BinaryOp::Add, points as u64);
let schema = evaluated_schema(grid);
let pool = bounded_pool(4 << 20);
let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new());
execution
.push_evaluated_at(
&left,
evaluated_chunk_of(&schema, metric(), 0, vec![Some(1.0); points], true),
)
.unwrap();
assert!(execution.next_chunk().unwrap().is_none());
// Both right chunks close the one label set, so its match group has two members.
for _ in 0..2 {
execution
.push_evaluated_at(
&right,
evaluated_chunk_of(&schema, metric(), 0, vec![None; points], true),
)
.unwrap();
assert!(execution.next_chunk().unwrap().is_none());
}
let finished = execution.finish();
assert!(
finished.is_ok(),
"an all-missing match group over {points} points does not fit a 4 MiB pool: {:?}",
finished.as_ref().err()
);
assert!(
finished.unwrap().is_empty(),
"no point holds a value on both sides, so nothing is answered"
);
assert_eq!(pool.reserved(), 0);
assert_eq!(execution.state().reserved_bytes(), 0);
}
+64 -34
View File
@@ -234,6 +234,9 @@ fn binary_vectors_inner(
merge_fields(fields, output, budget)
}
/// The capacity a fresh `Vec` holds: a first push allocates four entries, and a full one doubles.
const VEC_MIN_CAPACITY: usize = 4;
/// One label of a match key, charged twice because `group_key` grows its `Vec` by doubling.
const KEY_LABEL_BYTES: usize = 2 * std::mem::size_of::<(Arc<str>, Arc<str>)>();
@@ -374,10 +377,6 @@ impl<'a> Budget<'a> {
}
}
/// The slots the duplicate scan's point list holds before it doubles: a `Vec` of `usize` starts
/// with room for four, so a scan that collects a single point still allocates that much.
const POINT_LIST_FIRST_SLOTS: usize = 4;
/// What one field's tables, flags and messages cost before it pairs anything.
#[derive(Default)]
struct FieldMetadata {
@@ -400,7 +399,7 @@ struct FieldMetadata {
/// around one rendering. Nothing here counts the pairing the field walks: which series match, where
/// their axes overlap and which rows the operator keeps is only known as the walk reaches them.
/// The duplicate scan's list of present points is not part of these tables either: a group of one
/// member collects nothing, and a wider group's list is charged by [`overlap`] as it grows it.
/// member collects nothing, and a wider group's list is charged by [`overlap`] as it is pushed to.
fn field_metadata_bytes(
matching: &BinaryMatching,
include: &[Arc<str>],
@@ -543,13 +542,12 @@ fn result_set_bytes(metric: &MetricName, one: &MetricName, include: &[Arc<str>])
///
/// The merge files every kept triple under its label set and hands one series per label set to a
/// builder, so the triples, the table, the row it fills, the output vector and the builder's own
/// buffers are alive together: one triple of a label set's samples, one index per retained row, the
/// span a run may start with, a row of the histogram table per histogram and the heap its buckets
/// are copied to. Every term counts a value a field actually kept or a label set it filed, so a pair
/// the operator filtered costs its label set -- the table holds one without a value -- and no row at
/// all. The builder's value columns are not here: they hold one slot per output row and field, and a
/// row holding only some of the fields still writes a null into the others, so [`merge_fields`]
/// charges them per row as it pushes it.
/// buffers are alive together: one triple of a label set's samples, the span a run may start with, a
/// row of the histogram table per histogram and the heap its buckets are copied to. Every term counts
/// a value a field actually kept or a label set it filed, so a pair the operator filtered costs its
/// label set -- the table holds one without a value -- and no row at all. What grows with the output
/// rows is charged by [`merge_fields`] as they are pushed: the value and validity column of every
/// field, and the index column of a field whose first histogram backfills the rows before it.
fn merge_bytes(
kept: usize,
sets: usize,
@@ -558,18 +556,19 @@ fn merge_bytes(
histogram_heap: usize,
) -> usize {
// One kept triple, charged twice for the capacity its vector ends at, one slot of the row the
// builder reads, the span it may start and its row of the histogram table's index column.
// builder reads and the span it may start.
let row = (2 * std::mem::size_of::<(usize, usize, GridSample)>())
.saturating_add(2 * std::mem::size_of::<usize>())
.saturating_add(2 * std::mem::size_of::<GridSpan>())
.saturating_add(std::mem::size_of::<Option<u32>>());
.saturating_add(2 * std::mem::size_of::<GridSpan>());
kept.saturating_mul(row)
// One label set's entry in the merge table and the series it is written out as.
.saturating_add(
sets.saturating_mul(MERGE_SLOT_BYTES.saturating_add(2 * std::mem::size_of::<Series>())),
)
// The row the builder reads its pushes from, one slot per field.
// The row the builder reads its pushes from and the per-field ledger of the capacity each
// histogram index column has reached, one slot per field.
.saturating_add(2 * fields.saturating_mul(std::mem::size_of::<Option<GridSample>>()))
.saturating_add(fields.saturating_mul(std::mem::size_of::<usize>()))
// The histogram table of every field that carries one: the builders of its columns, the
// scalars and list offsets of its rows, and the buckets the batch copies beside the value
// the results still hold.
@@ -1107,11 +1106,25 @@ fn set_field(
Ok(results)
}
/// The capacity a field's histogram index column reaches when `pushed` output rows have been pushed
/// and the row being pushed holds `histogram`; `None` when the column does not grow.
///
/// [`GridSeriesBuilder::push_samples`] creates the column at the first histogram with one entry per
/// row before it and then pushes, so that push allocates `max(2 * pushed, 4)`; every later row adds
/// an entry, and a full column doubles.
fn index_column_growth(pushed: usize, capacity: usize, histogram: bool) -> Option<usize> {
if capacity == 0 {
return histogram.then(|| (2 * pushed).max(VEC_MIN_CAPACITY));
}
(pushed == capacity).then(|| 2 * capacity)
}
/// Joins the per-field results into series with all fields, by label set.
///
/// The builder's value columns grow with the output rows, which only exist once the triples are
/// being grouped, so they are reserved per row -- one row's worth before the builder takes it --
/// when a caller accounts for the workspace.
/// The builder's columns grow with the output rows, which only exist once the triples are being
/// grouped, so they are reserved per row -- one row's worth before the builder takes it -- when a
/// caller accounts for the workspace: the value and validity column of every field, and the
/// histogram index column of a field that holds one.
fn merge_fields(
fields: Vec<Results>,
schema: Arc<SeriesSchema>,
@@ -1139,8 +1152,13 @@ fn merge_fields(
for (metric, mut samples) in merged.into_values() {
// One builder serves every label set, so each result series starts explicitly.
builder.begin_series();
// The capacity each field's histogram index column has reached, `0` until its first
// histogram. The columns are handed out inside the series, so what they are charged stays
// charged; only the growth of one series is charged at a time.
let mut index_capacity = vec![0usize; width];
samples.sort_unstable_by_key(|(point, field, _)| (*point, *field));
let mut samples = samples.into_iter().peekable();
let mut pushed = 0usize;
while let Some((point, field, value)) = samples.next() {
row.clear();
row.resize_with(width, || None);
@@ -1151,8 +1169,22 @@ fn merge_fields(
// The columns of this row, admitted before the builder grows them.
if let Some(budget) = budget.as_deref_mut() {
budget.retain(merge_row_bytes(width))?;
let mut index = 0usize;
for (capacity, value) in index_capacity.iter_mut().zip(&row) {
let histogram = matches!(value, Some(GridSample::Histogram(_)));
if let Some(grown) = index_column_growth(pushed, *capacity, histogram) {
index = index.saturating_add(
grown
.saturating_sub(*capacity)
.saturating_mul(std::mem::size_of::<Option<u32>>()),
);
*capacity = grown;
}
}
budget.retain(index)?;
}
builder.push_samples(point, &row)?;
pushed += 1;
}
output.extend(builder.finish(metric)?);
}
@@ -1223,12 +1255,10 @@ fn group(operand: &Operand<'_>, labels: &MatchingLabels) -> HashMap<GroupKey, Ve
/// The first point where two `members` are present and `relevant` holds.
///
/// The one list of the group's present points is the whole of what this scan allocates, and it is
/// the field's workspace beside its `tables` rather than part of them: it is reserved before a point
/// is pushed -- one slot per row the group's members carry at most, never fewer than the four a new
/// list starts with, doubled because the list grows by doubling -- and given back once the list is
/// dropped. A group of one member cannot repeat a point, so it allocates nothing and is charged
/// nothing.
/// The list of the group's present points is the only thing this scan allocates, and it is charged
/// as a slot is pushed: the figure follows the list's capacity, so a group of one member and a group
/// whose members carry no sample both cost nothing, however long their axes are. The list is given
/// back once it is dropped.
fn overlap(
operand: &Operand<'_>,
members: &[usize],
@@ -1239,15 +1269,6 @@ fn overlap(
if members.len() < 2 {
return Ok(None);
}
if let Some(budget) = budget.as_deref_mut() {
let rows = members
.iter()
.fold(0usize, |rows, &member| {
rows.saturating_add(operand.axis(member).len())
})
.max(POINT_LIST_FIRST_SLOTS);
budget.step(tables.saturating_add(2 * std::mem::size_of::<usize>() * rows))?;
}
// Sorting the present points keeps the cost proportional to them, not to the span.
let mut points = Vec::new();
for &member in members {
@@ -1257,6 +1278,15 @@ fn overlap(
cursor.at(row).expect("a row never goes back"),
SampleRef::Missing
) {
// The push that starts a new allocation: the capacity the list reaches is what is
// live beside the field's tables, charged twice for the buffer it replaces.
if let Some(budget) = budget.as_deref_mut()
&& points.len() == points.capacity()
{
let capacity = points.capacity().saturating_mul(2).max(VEC_MIN_CAPACITY);
let list = capacity.saturating_mul(2 * std::mem::size_of::<usize>());
budget.step(tables.saturating_add(list))?;
}
points.push(point);
}
}
+203
View File
@@ -2854,3 +2854,206 @@ fn fields_holding_disjoint_points_are_charged_their_columns() {
"peak {peak} for {points} rows of {fields} fields"
);
}
/// A duplicate scan is charged the points it collects: a group whose members carry no sample
/// collects nothing, however long their axes are, while the points a group does carry are its own
/// slots.
#[test]
fn a_duplicate_scan_is_charged_the_points_it_collects() {
let points = 100_000;
let grid = budget_grid(points);
let schema = Arc::new(schema(1, &grid));
// A member of the group with no value at any point: its axis is the whole grid.
let null_member = |name: &str| -> Series {
Series {
metric: metric(&Labels::from([
("__name__".to_string(), name.to_string()),
("job".to_string(), "x".to_string()),
])),
time: TimeAxis::Grid(GridRange::try_new(grid.clone(), 0, points).unwrap()),
fields: SeriesFields::one(Samples::from_float(Float64Array::new_null(points))),
}
};
let sparse_member = |name: &str, at: &[usize]| -> Series {
let labels = Labels::from([
("__name__".to_string(), name.to_string()),
("job".to_string(), "x".to_string()),
]);
let values: Points = at
.iter()
.map(|&position| (position, f64::to_bits(position as f64)))
.collect();
sparse_series(&schema, &labels, &values).expect("the operand holds a value")
};
// The left side carries one value, at a point no member carries, so the group is never a
// duplicate and no row is filed.
let lhs = vec![sparse_member("a", &[90_000])];
let matching = budget_matching(BinaryOp::Add, MatchingLabels::On(vec![Arc::from("job")]));
// What the run holds beside the list: three presence flags per grid point.
let floor = 3 * points;
let run = |rhs: &[Series], limit: usize| -> Result<Vec<Series>, DataFusionError> {
budgeted(&matching, &schema, &lhs, &schema, rhs, limit).0
};
// Two members of one group, each with a value at no point: their axes are the whole grid and
// the list stays empty.
let empty_group = vec![null_member("b"), null_member("c")];
assert!(
binary_vectors(&matching, &schema, &lhs, &schema, &empty_group)
.expect("no member carries a sample")
.is_empty()
);
let output = run(&empty_group, floor + 64 * 1024).expect("an empty list costs nothing");
assert!(output.is_empty(), "{output:?}");
// A few points of each member are the list's own slots, and still fit beside the flags.
let few = vec![
sparse_member("b", &[10, 20, 30, 40]),
sparse_member("c", &[50, 60, 70, 80]),
];
assert!(
run(&few, floor + 64 * 1024)
.expect("a few points cost a few slots")
.is_empty()
);
// Thousands of points are charged as they are collected: the same limit refuses them, since
// their slots alone are tens of kilobytes.
let many = vec![
sparse_member("b", &(0..4_000).collect::<Vec<_>>()),
sparse_member("c", &(4_000..8_000).collect::<Vec<_>>()),
];
let error = run(&many, floor + 64 * 1024)
.expect_err("the collected points are charged as they are pushed");
assert!(
matches!(error, DataFusionError::ResourcesExhausted(_)),
"{error}"
);
}
/// A field's histogram index column is charged as it grows: the builder creates it over the rows
/// before the field's first histogram and grows it with every row after, so a wide result whose
/// fields meet their histograms late holds those columns beside its value columns.
#[test]
fn sparse_wide_results_are_charged_their_index_columns() {
let points = 1_025;
let fields = 64;
let grid = budget_grid(points);
let schema = Arc::new(schema(fields, &grid));
// The value of a field at its one histogram point.
let histogram = || {
GridSample::Histogram(Box::new(
FloatHistogram::try_from_view(&OwnedHistogram::parse("0/0/0/0/0/0/////").view())
.expect("an empty histogram is readable"),
))
};
// Field `f` holds a value only at the points `f` apart: a float at every one of them but the
// last, and an empty histogram at its single point of the last `fields` rows. Every field meets
// its first histogram late and keeps the column to the end.
let side = |name: &str, offset: f64| -> Vec<Series> {
let mut builder = GridSeriesBuilder::try_new(schema.clone()).unwrap();
builder.begin_series();
for point in 0..points {
let mut row: Vec<Option<GridSample>> = Vec::with_capacity(fields);
row.resize_with(fields, || None);
row[point % fields] = Some(if point >= points - fields {
histogram()
} else {
GridSample::Float(point as f64 + offset)
});
builder.push_samples(point, &row).unwrap();
}
builder
.finish(metric(&Labels::from([
("__name__".to_string(), name.to_string()),
("job".to_string(), "x".to_string()),
])))
.unwrap()
.into_iter()
.collect()
};
let matching = budget_matching(BinaryOp::Add, MatchingLabels::Ignoring(Vec::new()));
let (output, seen) = budgeted(
&matching,
&schema,
&side("a", 1.0),
&schema,
&side("b", 2.0),
usize::MAX,
);
let output = output.expect("every field adds up");
assert_eq!(output.len(), 1, "{output:?}");
// Every point keeps the field it belongs to and nothing in the others, the floats add up and
// the empty histograms still hold an empty histogram.
let mut cursors: Vec<RowCursor<'_>> = output[0]
.fields
.iter()
.map(|field| field.cursor())
.collect();
let mut rows = 0;
for (point, row) in output[0].time.positions() {
rows += 1;
for (field, cursor) in cursors.iter_mut().enumerate() {
let owned = field == point % fields;
match cursor.at(row).expect("a row never goes back") {
SampleRef::Float(value) => {
assert!(owned, "point {point} keeps field {field}");
assert!(point < points - fields, "point {point} is not a histogram");
assert_eq!(value, 2.0 * point as f64 + 3.0, "point {point}");
}
SampleRef::Histogram(view) => {
assert!(owned, "point {point} keeps field {field}");
assert!(point >= points - fields, "point {point} is a histogram");
assert_eq!(view.count, 0.0, "empty histograms add up empty");
assert!(
view.positive.buckets.is_empty() && view.negative.buckets.is_empty(),
"point {point} keeps an empty histogram"
);
}
SampleRef::Missing => assert!(!owned, "point {point} holds field {field}"),
}
}
}
assert_eq!(rows, points, "every point is a row");
// The capacity Rust gives these columns for this fixture, from its growth rule alone: a push
// into a full list allocates at least four entries and otherwise doubles. One value column per
// field, pushed once per row, and one index column per field, created over the rows before its
// first histogram and pushed once per row after.
let grow = |capacity: usize, len: usize| -> usize {
if len < capacity {
capacity
} else {
(2 * capacity).max(4)
}
};
let mut values = 0usize;
let mut index = vec![0usize; fields];
for pushed in 0..points {
values = grow(values, pushed);
let late = pushed >= points - fields;
for (field, capacity) in index.iter_mut().enumerate() {
if *capacity == 0 {
if late && pushed % fields == field {
// Created over the `pushed` rows before the field's first histogram, then
// pushed to, so its length is `pushed + 1` like every column made before it.
*capacity = grow(pushed, pushed);
}
} else {
*capacity = grow(*capacity, pushed);
}
}
}
let columns = fields * values * std::mem::size_of::<f64>()
+ index.iter().fold(0usize, |bytes, capacity| {
bytes.saturating_add(capacity.saturating_mul(std::mem::size_of::<Option<u32>>()))
});
let peak = *seen
.iter()
.max()
.expect("the operation asks for its workspace");
assert!(
peak >= columns,
"peak {peak} below the {columns} bytes of value and index columns the result holds"
);
assert!(
peak <= 2 * columns,
"peak {peak} is more than the {columns} bytes of columns the result holds"
);
}