From 4645ab9ed74c9433bbc0a99a1208243c9df7019f Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 20 Oct 2026 03:41:18 +0000 Subject: [PATCH] perf(mito2): gather last rows with one take per input batch (#9737) * 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 * test(mito2): cover duplicate last timestamps in the take path Signed-off-by: Dennis Zhuang * 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 --------- Signed-off-by: Dennis Zhuang --- src/mito2/AGENTS.md | 3 + src/mito2/src/read/last_row.rs | 167 +++++++++++++++++++++++++++------ 2 files changed, 140 insertions(+), 30 deletions(-) diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index c41aa649111..811632b1e26 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -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 diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index 2bafe02dbbe..b7da3e1f70c 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -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::() + .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::(); + 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();