rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-21 19:32:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(mito2): gather last rows with one take per input batch (#9737)

Browse Source 
* perf(mito2): gather last rows with one take per input batch

FlatLastTimestampSelector buffered a single-row slice per series and
concatenated up to 8192 of them. Concatenating slices with dictionary columns
merges the dictionaries once per slice, which made high-cardinality instant
PromQL queries using the LastRow scan path (#9034) slower. Select the rows of
series that end inside a batch with one take instead.

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* test(mito2): cover duplicate last timestamps in the take path

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(mito2): group last rows by primary key bytes

Flat batches may assign different dictionary keys to the same primary key, for
example when binary keys are wrapped with identity keys. Merge adjacent ranges
with equal key bytes before selecting last rows.

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2026-10-20 03:41:18 +00:00
committed by GitHub
parent 9c7e1e1605
commit 4645ab9ed7
2 changed files with 140 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/AGENTS.md
+3
View File
@@ -98,6 +98,9 @@ Tests live next to the code as `*_test.rs` (e.g. `src/mito2/src/engine/flush_tes
only close/drop unregisters them. See `compaction/scheduler/state.rs`.
- Local compaction units use the snapshot-free `CompactionContext`. Retaining a
full `CompactionRegion` in a unit pins other units' inputs and delays SST purge.
- In flat batches, equal `__primary_key` dictionary keys imply the same series, but
not the reverse: `wrap_pk_binary_to_dict` gives every row its own key. Compare key
bytes when grouping rows by series.
## Maintenance contract
src/mito2/src/read/last_row.rs
+137 -30
View File
@@ -16,8 +16,8 @@
use std::sync::Arc;
use datatypes::arrow::array::{Array, BinaryArray};
use datatypes::arrow::compute::concat_batches;
use datatypes::arrow::array::{Array, BinaryArray, UInt32Array};
use datatypes::arrow::compute::{concat_batches, take_record_batch};
use datatypes::arrow::record_batch::RecordBatch;
use futures::{Stream, TryStreamExt};
use snafu::ResultExt;
@@ -155,6 +155,11 @@ impl BatchBuffer {
self.batches.extend_from_slice(batches);
}
fn push(&mut self, batch: RecordBatch) {
self.num_rows += batch.num_rows();
self.batches.push(batch);
}
/// Returns true if the buffer has no batches.
fn is_empty(&self) -> bool {
self.batches.is_empty()
@@ -306,41 +311,70 @@ impl FlatLastTimestampSelector {
return Ok(());
}
// Dictionary keys may differ for the same primary key, e.g. readers that wrap binary
// keys into a dictionary with identity keys, so merge adjacent ranges by key bytes.
let mut series_ranges: Vec<(usize, usize)> = Vec::with_capacity(offsets.len() - 1);
for window in offsets.windows(2) {
let (start, end) = (window[0], window[1]);
match series_ranges.last_mut() {
Some(last)
if primary_key_bytes_at(&batch, pk_col_idx, last.0)
== primary_key_bytes_at(&batch, pk_col_idx, start) =>
{
last.1 = end;
}
_ => series_ranges.push((start, end)),
}
}
let ts_values = timestamp_array_to_i64_slice(batch.column(ts_col_idx));
for i in 0..offsets.len() - 1 {
let range_start = offsets[i];
let range_end = offsets[i + 1];
let num_ranges = series_ranges.len();
// Rows of keys that start and end inside this batch. They are gathered with one
// `take` instead of buffering a slice per key: concatenating thousands of single-row
// slices merges their dictionaries once per slice.
let mut completed_rows = Vec::new();
for (i, &(range_start, range_end)) in series_ranges.iter().enumerate() {
let range_key = primary_key_bytes_at(&batch, pk_col_idx, range_start);
let range_last_ts = ts_values[range_end - 1];
let range_last_ts_start = last_timestamp_start(ts_values, range_start, range_end);
let range_slice = batch.slice(range_last_ts_start, range_end - range_last_ts_start);
let is_last_range = i + 1 == num_ranges;
match self.current_key.as_mut() {
Some(state) if state.key.as_slice() == range_key => {
if range_last_ts > state.last_timestamp {
state.last_timestamp = range_last_ts;
state.slices.clear();
state.slices.push(range_slice);
} else if range_last_ts == state.last_timestamp {
state.slices.push(range_slice);
}
// Only the first range can continue the key carried over from previous batches.
if let Some(state) = self.current_key.as_mut()
&& state.key.as_slice() == range_key
{
let range_slice = batch.slice(range_last_ts_start, range_end - range_last_ts_start);
if range_last_ts > state.last_timestamp {
state.last_timestamp = range_last_ts;
state.slices.clear();
state.slices.push(range_slice);
} else if range_last_ts == state.last_timestamp {
state.slices.push(range_slice);
}
Some(_) => {
if !is_last_range {
self.flush_current_key(output_buffer);
self.current_key = Some(LastKeyState::new(
range_key.to_vec(),
range_last_ts,
range_slice,
));
}
None => {
self.current_key = Some(LastKeyState::new(
range_key.to_vec(),
range_last_ts,
range_slice,
));
}
continue;
}
// Flushes the carried-over key before rows of this batch to keep the output sorted.
self.flush_current_key(output_buffer);
if is_last_range {
// The key may continue in the next batch.
self.current_key = Some(LastKeyState::new(
range_key.to_vec(),
range_last_ts,
batch.slice(range_last_ts_start, range_end - range_last_ts_start),
));
} else {
completed_rows.extend(range_last_ts_start as u32..range_end as u32);
}
}
if !completed_rows.is_empty() {
let indices = UInt32Array::from(completed_rows);
let selected = take_record_batch(&batch, &indices).context(ComputeArrowSnafu)?;
output_buffer.push(selected);
}
Ok(())
@@ -429,8 +463,8 @@ mod tests {
use std::sync::Arc;
use datatypes::arrow::array::{
ArrayRef, BinaryDictionaryBuilder, Int64Array, TimestampMillisecondArray, UInt8Array,
UInt64Array,
ArrayRef, BinaryDictionaryBuilder, DictionaryArray, Int64Array, TimestampMillisecondArray,
UInt8Array, UInt64Array,
};
use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit, UInt32Type};
use datatypes::arrow::record_batch::RecordBatch;
@@ -630,6 +664,79 @@ mod tests {
assert_eq!(vec![(b"k1".to_vec(), 4), (b"k1".to_vec(), 4)], results);
}
#[test]
fn test_flat_concat_across_parent_dictionaries() {
let mut selector = FlatLastTimestampSelector::default();
let mut output_buffer = BatchBuffer::new();
let batches = vec![
new_flat_batch(&[b"a", b"b", b"b", b"c"], &[1, 1, 2, 3], &[10, 20, 21, 30]),
new_flat_batch(
&[b"c", b"d", b"d", b"e", b"e"],
&[3, 4, 4, 5, 5],
&[31, 40, 41, 50, 51],
),
new_flat_batch(&[b"f", b"g"], &[6, 7], &[60, 70]),
];
for batch in batches {
selector.on_next(batch, &mut output_buffer).unwrap();
}
selector.finish(&mut output_buffer).unwrap();
let merged = output_buffer.concat().unwrap();
let mut rows = Vec::new();
extract_flat_rows(&merged, &mut rows);
let fields = merged
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.to_vec();
assert_eq!(
vec![
(b"a".to_vec(), 1),
(b"b".to_vec(), 2),
(b"c".to_vec(), 3),
(b"c".to_vec(), 3),
(b"d".to_vec(), 4),
(b"d".to_vec(), 4),
(b"e".to_vec(), 5),
(b"e".to_vec(), 5),
(b"f".to_vec(), 6),
(b"g".to_vec(), 7),
],
rows
);
assert_eq!(vec![10, 21, 30, 31, 40, 41, 50, 51, 60, 70], fields);
}
/// Rebuilds the primary key column with one dictionary key per row, as
/// `wrap_pk_binary_to_dict` does for binary primary keys.
fn with_identity_pk_keys(batch: RecordBatch) -> RecordBatch {
let pk_col_idx = primary_key_column_index(batch.num_columns());
let values = (0..batch.num_rows())
.map(|i| Some(primary_key_bytes_at(&batch, pk_col_idx, i)))
.collect::<BinaryArray>();
let keys = UInt32Array::from_iter_values(0..batch.num_rows() as u32);
let mut columns = batch.columns().to_vec();
columns[pk_col_idx] = Arc::new(DictionaryArray::new(keys, Arc::new(values)));
RecordBatch::try_new(batch.schema(), columns).unwrap()
}
#[test]
fn test_flat_identity_dictionary_keys() {
let mut selector = FlatLastTimestampSelector::default();
let batches = vec![
with_identity_pk_keys(new_flat_batch(&[b"a", b"a", b"b"], &[1, 2, 3], &[0; 3])),
with_identity_pk_keys(new_flat_batch(&[b"b", b"c", b"c"], &[4, 5, 6], &[0; 3])),
];
let results = collect_flat_results(&mut selector, batches);
assert_eq!(
vec![(b"a".to_vec(), 2), (b"b".to_vec(), 4), (b"c".to_vec(), 6)],
results
);
}
#[test]
fn test_flat_finish_is_one_shot() {
let mut selector = FlatLastTimestampSelector::default();