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();