diff --git a/Cargo.lock b/Cargo.lock index f2b9ae85d0..fb277f0b0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8705,6 +8705,7 @@ dependencies = [ "toml 0.8.23", "tracing", "uuid", + "winner_tree", "zstd", ] @@ -17129,6 +17130,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "winner_tree" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0639a3cbe35d6df7643f015e6fbde7664c7daf6d376b802b3c0fa05141b316e7" + [[package]] name = "winnow" version = "0.5.40" diff --git a/Cargo.toml b/Cargo.toml index 77fd4dbade..55181077f8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -273,6 +273,7 @@ uddsketch-rs = "0.1.0" uuid = { version = "1.17", features = ["serde", "v4", "v7", "fast-rng"] } vrl = "0.33" vu128 = "1.1" +winner_tree = "0.1.1" zstd = "0.13" # DO_NOT_REMOVE_THIS: END_OF_EXTERNAL_DEPENDENCIES diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 116c4066ae..100a6fb102 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -93,6 +93,7 @@ tokio-stream.workspace = true tokio-util.workspace = true tracing.workspace = true uuid.workspace = true +winner_tree.workspace = true zstd.workspace = true [dev-dependencies] diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 6d90f90b78..edaa97ff58 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -284,6 +284,92 @@ async fn test_write_query_region_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +/// Builds rows with `tag_num` string tags, one float field and a timestamp +/// in second resolution, like [build_rows] but for various tag counts. +fn build_rows_with_tags(tag_num: usize, rows: &[(usize, i64)]) -> Vec { + rows.iter() + .map(|&(value, ts)| { + let mut values: Vec = (0..tag_num) + .map(|_| ValueData::StringValue(value.to_string())) + .collect(); + values.push(ValueData::F64Value(value as f64)); + values.push(ValueData::TimestampMillisecondValue(ts * 1000)); + row(values) + }) + .collect() +} + +#[tokio::test] +async fn test_flat_merge_with_various_schemas() { + // Tables without tags: all rows share an empty primary key. + let expected_no_tag = "\ ++---------+---------------------+ +| field_0 | ts | ++---------+---------------------+ +| 0.0 | 1970-01-01T00:00:00 | +| 1.0 | 1970-01-01T00:00:01 | +| 2.0 | 1970-01-01T00:00:02 | +| 3.0 | 1970-01-01T00:00:03 | ++---------+---------------------+"; + flat_merge_schema_check(0, "flat-merge-no-tag", expected_no_tag).await; + + // Tables with multiple string tags: tag columns are dictionary-encoded + // in the flat input schema of the merge reader. + let expected_multi_tag = "\ ++-------+-------+---------+---------------------+ +| tag_0 | tag_1 | field_0 | ts | ++-------+-------+---------+---------------------+ +| 0 | 0 | 0.0 | 1970-01-01T00:00:00 | +| 1 | 1 | 1.0 | 1970-01-01T00:00:01 | +| 2 | 2 | 2.0 | 1970-01-01T00:00:02 | +| 3 | 3 | 3.0 | 1970-01-01T00:00:03 | ++-------+-------+---------+---------------------+"; + flat_merge_schema_check(2, "flat-merge-multi-tag", expected_multi_tag).await; +} + +/// Writes and flushes rows, then writes interleaved rows so a scan must merge +/// the SST with the memtable through the flat merge reader, and checks the +/// merged output. +async fn flat_merge_schema_check(tag_num: usize, prefix: &str, expected: &str) { + let mut env = TestEnv::with_prefix(prefix).await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().tag_num(tag_num).build(); + + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Writes rows with even timestamps and flushes them to an SST. + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows_with_tags(tag_num, &[(0, 0), (2, 2)]), + }; + put_rows(&engine, region_id, rows).await; + flush_region(&engine, region_id, None).await; + + // Writes rows with odd timestamps; they stay in the memtable and + // interleave with the SST rows when merged. + let rows = Rows { + schema: column_schemas, + rows: build_rows_with_tags(tag_num, &[(1, 1), (3, 3)]), + }; + put_rows(&engine, region_id, rows).await; + + let request = ScanRequest::default(); + let stream = engine.scan_to_stream(region_id, request).await.unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(expected, batches.pretty_print().unwrap()); +} + #[tokio::test] async fn test_different_order() { test_different_order_with_format(false).await; diff --git a/src/mito2/src/read/flat_merge.rs b/src/mito2/src/read/flat_merge.rs index f1c30dfbfd..db368aedff 100644 --- a/src/mito2/src/read/flat_merge.rs +++ b/src/mito2/src/read/flat_merge.rs @@ -33,18 +33,19 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::arrow_array::BinaryArray; use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, TryStreamExt}; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::storage::SequenceNumber; -use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; +use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; +use winner_tree::WinnerTree; -use crate::error::{ComputeArrowSnafu, Result}; +use crate::error::{ComputeArrowSnafu, InvalidRecordBatchSnafu, Result}; use crate::memtable::BoxedRecordBatchIterator; use crate::metrics::READ_STAGE_ELAPSED; use crate::read::BoxedRecordBatchStream; use crate::sst::parquet::flat_format::{ primary_key_column_index, sequence_column_index, time_index_column_index, }; -use crate::sst::parquet::format::PrimaryKeyArray; +use crate::sst::parquet::format::{FIXED_POS_COLUMN_NUM, PrimaryKeyArray}; /// Checks whether interleaving the selected rows from byte columns would overflow /// i32 offsets. Similar to arrow-rs `interleave_bytes()`, accumulates offsets and @@ -305,6 +306,10 @@ pub struct BatchBuilder { impl BatchBuilder { /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size` pub fn new(schema: SchemaRef, stream_count: usize, batch_size: usize) -> Self { + // A real flat-format schema always has at least 4 columns (time index + // plus the 3 internal columns); the `>= 3` check below only keeps + // `primary_key_column_index` (`num_columns - 3`) from underflowing on + // generic schemas without internal columns. let primary_key_column_idx = (schema.fields.len() >= 3) .then(|| primary_key_column_index(schema.fields.len())) .filter(|&column_idx| schema.field(column_idx).name() == PRIMARY_KEY_COLUMN_NAME); @@ -430,199 +435,6 @@ impl BatchBuilder { } } -/// Sentinel for an empty slot in a [TournamentTree]. -const EMPTY_SLOT: usize = usize::MAX; - -/// A tournament tree over a fixed set of slots. -/// -/// Each node occupies one slot; a slot is empty while its node is not in the -/// tree (i.e. it is in the cold heap or reached EOF). An empty slot always -/// loses a match. The number of leaves is padded to a power of two; leaves -/// beyond the capacity never participate, so the tree works for arbitrary -/// capacities. -/// -/// Invariant: every internal tree node caches the champion (hottest slot, per -/// `Ord`) of its subtree, so the root always holds the hottest occupied slot -/// and updating a leaf only requires recomputing the ~log2(capacity) internal -/// nodes on its path to the root. We cache champions ("winner tree") so that -/// each internal node is a pure function of its children — insertion, removal -/// and mutation share one replay path — and the runner-up is directly -/// available from the sibling champions on the winner's path, which the -/// hot/cold transition check needs after mutating the winner in place. -struct TournamentTree { - /// Slot storage, one slot per node. `None` means the slot is empty. - nodes: Vec>, - /// Tree of champions: `tree[i]` is the slot index of the champion of the - /// subtree rooted at `i` (or [EMPTY_SLOT]). Leaves for slot `s` live at - /// index `leaves + s`; the root is index 1. - tree: Vec, - /// Number of occupied slots. - len: usize, - /// Number of leaves, padded to a power of two. - leaves: usize, - /// Cached `(winner slot, second-best slot)`, with [EMPTY_SLOT] standing - /// for no second-best. Invalidated by [TournamentTree::replay], the single - /// choke point of every structural change (push, pop, winner replay). - /// - /// While the tree is structurally unchanged and the winner keeps its slot - /// the cache stays valid: only the winner's node may be mutated in place - /// and the second-best slot is never the winner's slot, so the cached - /// slot's node is untouched. - second_best_cache: Option<(usize, usize)>, -} - -impl TournamentTree { - fn with_capacity(capacity: usize) -> Self { - let leaves = capacity.next_power_of_two().max(1); - let mut nodes = Vec::new(); - nodes.resize_with(capacity, || None); - Self { - nodes, - tree: vec![EMPTY_SLOT; leaves * 2], - len: 0, - leaves, - second_best_cache: None, - } - } - - fn is_empty(&self) -> bool { - self.len == 0 - } - - fn len(&self) -> usize { - self.len - } - - /// Returns the winner (greatest element among occupied slots). - fn peek(&self) -> Option<&T> { - self.winner_slot() - .and_then(|slot| self.nodes[slot].as_ref()) - } - - /// Returns the winner mutably. Call [TournamentTree::replay_winner] after - /// mutating it, or remove it with `pop()` if it reached EOF or moved cold. - fn winner_mut(&mut self) -> Option<&mut T> { - let slot = self.winner_slot()?; - self.nodes[slot].as_mut() - } - - /// Returns the second greatest element among occupied slots. - /// - /// The runner-up is the champion of one of the sibling subtrees on the - /// winner's path to the root. - #[cfg(test)] - fn second_best(&mut self) -> Option<&T> { - let slot = self.second_best_slot()?; - Some(self.nodes[slot].as_ref().unwrap()) - } - - /// Returns the winner and the second greatest element among occupied slots. - fn winner_and_second_best(&mut self) -> (Option<&T>, Option<&T>) { - let second = self.second_best_slot(); - let winner = self.winner_slot(); - ( - winner.map(|slot| self.nodes[slot].as_ref().unwrap()), - second.map(|slot| self.nodes[slot].as_ref().unwrap()), - ) - } - - /// Inserts `value` into a free slot and replays its path to the root. - /// - /// Scans for a free slot instead of keeping a free list to keep the tree - /// cheap to construct. This is O(capacity), but pushes only happen on - /// batch transitions, never per row. - /// - /// # Panics - /// Panics if the tree is full. - fn push(&mut self, value: T) { - let slot = self - .nodes - .iter() - .position(Option::is_none) - .expect("tournament tree is full"); - self.nodes[slot] = Some(value); - self.tree[self.leaves + slot] = slot; - self.len += 1; - self.replay(slot); - } - - /// Removes and returns the winner, leaving its slot empty. - fn pop(&mut self) -> Option { - let slot = self.winner_slot()?; - let value = self.nodes[slot].take(); - debug_assert!(value.is_some()); - self.tree[self.leaves + slot] = EMPTY_SLOT; - self.len -= 1; - self.replay(slot); - value - } - - /// Replays the winner's path after its value was mutated in place. - fn replay_winner(&mut self) { - if let Some(slot) = self.winner_slot() { - self.replay(slot); - } - } - - /// Returns the slot of the champion at the root, if any slot is occupied. - fn winner_slot(&self) -> Option { - (self.tree[1] != EMPTY_SLOT).then_some(self.tree[1]) - } - - /// Returns the slot of the second greatest element among occupied slots. - fn second_best_slot(&mut self) -> Option { - let winner = self.winner_slot()?; - if let Some((cached_winner, cached)) = self.second_best_cache - && cached_winner == winner - { - return (cached != EMPTY_SLOT).then_some(cached); - } - let second = self.compute_second_best_slot(winner).unwrap_or(EMPTY_SLOT); - self.second_best_cache = Some((winner, second)); - (second != EMPTY_SLOT).then_some(second) - } - - /// Scans the champions of the sibling subtrees on `winner`'s path to the - /// root for the hottest one. - fn compute_second_best_slot(&self, winner: usize) -> Option { - let mut node = self.leaves + winner; - let mut best = EMPTY_SLOT; - while node > 1 { - let challenger = self.tree[node ^ 1]; - if self.wins(challenger, best) { - best = challenger; - } - node /= 2; - } - (best != EMPTY_SLOT).then_some(best) - } - - /// Returns true if slot `a` wins its match against slot `b`, i.e. `a` is - /// the greater element. An occupied slot always beats an empty slot - /// ([EMPTY_SLOT] or a slot whose node was removed). - fn wins(&self, a: usize, b: usize) -> bool { - match (self.nodes.get(a), self.nodes.get(b)) { - (Some(Some(x)), Some(Some(y))) => x >= y, - (Some(Some(_)), _) => true, - _ => false, - } - } - - /// Recomputes the internal nodes on the path from `slot`'s leaf to the - /// root (~log2(capacity) comparisons). - fn replay(&mut self, slot: usize) { - debug_assert!(slot < self.nodes.len()); - self.second_best_cache = None; - let mut node = (self.leaves + slot) / 2; - while node > 0 { - let left = self.tree[node * 2]; - let right = self.tree[node * 2 + 1]; - self.tree[node] = if self.wins(left, right) { left } else { right }; - node /= 2; - } - } -} - /// A comparable node of the heap. trait NodeCmp: Eq + Ord { /// Returns whether the node still has batch to read. @@ -643,7 +455,7 @@ struct MergeAlgo { /// /// Node in this tree **MUST** not be empty. A `merge window` is the (primary key, timestamp) /// range of the **winner node** in the `hot` tree. - hot: TournamentTree, + hot: WinnerTree, /// Holds nodes whose key range of current batch **isn't** overlapped with the merge window. /// /// Nodes in this heap **MUST** not be empty. @@ -657,7 +469,7 @@ impl MergeAlgo { fn new(mut nodes: Vec) -> Self { // Skips EOF nodes. nodes.retain(|node| !node.is_eof()); - let hot = TournamentTree::with_capacity(nodes.len()); + let hot = WinnerTree::with_capacity(nodes.len()); let cold = BinaryHeap::from(nodes); let mut algo = MergeAlgo { hot, cold }; @@ -763,41 +575,77 @@ struct SortColumns { } impl SortColumns { - /// Creates a new [SortColumns] from a [RecordBatch] and the position of the time index column. + /// Creates a new [SortColumns] from a [RecordBatch] in the flat format. /// - /// # Panics - /// Panics if the input batch doesn't have correct internal columns. - fn new(batch: &RecordBatch) -> Self { + /// Returns an error if the batch doesn't carry the flat-format internal + /// columns (time index, `__primary_key`, `__sequence`) of the expected + /// types at the fixed trailing positions. Unlike [BatchBuilder], which + /// falls back to plain `interleave` on generic schemas, row comparison + /// fundamentally requires these columns, so a batch without them is + /// rejected with an error instead of a panic. + fn try_new(batch: &RecordBatch) -> Result { let num_columns = batch.num_columns(); + ensure!( + num_columns >= FIXED_POS_COLUMN_NUM, + InvalidRecordBatchSnafu { + reason: format!( + "flat merge batch only has {num_columns} columns, expect at least {FIXED_POS_COLUMN_NUM}" + ), + } + ); let primary_key = batch .column(primary_key_column_index(num_columns)) .as_any() .downcast_ref::() - .unwrap() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected a {PRIMARY_KEY_COLUMN_NAME} column of type Dictionary(UInt32, Binary) at index {}", + primary_key_column_index(num_columns), + ), + })? .clone(); let primary_key_values = primary_key .values() .as_any() .downcast_ref::() - .unwrap() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected Binary {PRIMARY_KEY_COLUMN_NAME} dictionary values, got {}", + primary_key.values().data_type() + ), + })? .clone(); let timestamp = batch.column(time_index_column_index(num_columns)); - let (timestamp, _unit) = timestamp_array_to_primitive(timestamp).unwrap(); + let (timestamp, _unit) = + timestamp_array_to_primitive(timestamp).with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected a timestamp time index column at index {}, got {}", + time_index_column_index(num_columns), + batch + .column(time_index_column_index(num_columns)) + .data_type(), + ), + })?; let sequence = batch .column(sequence_column_index(num_columns)) .as_any() .downcast_ref::() - .unwrap() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected a UInt64 {SEQUENCE_COLUMN_NAME} column at index {}", + sequence_column_index(num_columns), + ), + })? .clone(); - Self { + Ok(Self { primary_key, primary_key_values, timestamp, sequence, #[cfg(test)] primary_key_lookups: Cell::new(0), - } + }) } fn primary_key_at(&self, index: usize) -> &[u8] { @@ -923,6 +771,12 @@ impl Ord for RowCursor { /// Iterator to merge multiple sorted iterators into a single sorted iterator. /// /// All iterators must be sorted by primary key, time index, sequence desc. +/// +/// Input batches must be in the flat format: the last four columns are time +/// index, `__primary_key`, `__sequence` and `__op_type`. Ordering uses only +/// (primary key, time index, sequence desc); `__op_type` is required for +/// downstream flat dedup, but is not part of the ordering key. +/// The name-based gate in [BatchBuilder] only makes output assembly degrade gracefully on generic schemas, not sorting. pub struct FlatMergeIterator { /// The merge algorithm to maintain heaps. algo: MergeAlgo, @@ -1059,6 +913,12 @@ impl Iterator for FlatMergeIterator { /// Iterator to merge multiple sorted iterators into a single sorted iterator. /// /// All iterators must be sorted by primary key, time index, sequence desc. +/// +/// Input batches must be in the flat format: the last four columns are time +/// index, `__primary_key`, `__sequence` and `__op_type`. Row comparison +/// decodes these internal columns and returns an error on batches that don't +/// match the flat format; the name-based gate in [BatchBuilder] only makes +/// output assembly degrade gracefully on generic schemas, not sorting. pub struct FlatMergeReader { /// The merge algorithm to maintain heaps. algo: MergeAlgo, @@ -1316,7 +1176,7 @@ impl GenericNode { /// Returns the fetched new batch. fn advance_batch(&mut self) -> Result> { let batch = self.advance_inner_iter()?; - let columns = batch.as_ref().map(SortColumns::new); + let columns = batch.as_ref().map(SortColumns::try_new).transpose()?; self.cursor = columns.map(RowCursor::new); Ok(batch) @@ -1355,7 +1215,7 @@ impl GenericNode { /// Returns the fetched new batch. async fn advance_batch(&mut self) -> Result> { let batch = self.advance_inner_iter().await?; - let columns = batch.as_ref().map(SortColumns::new); + let columns = batch.as_ref().map(SortColumns::try_new).transpose()?; self.cursor = columns.map(RowCursor::new); Ok(batch) @@ -1395,7 +1255,10 @@ mod tests { use api::v1::OpType; use datatypes::arrow::array::builder::BinaryDictionaryBuilder; - use datatypes::arrow::array::{Int64Array, TimestampMillisecondArray, UInt8Array, UInt64Array}; + use datatypes::arrow::array::{ + DictionaryArray, Int64Array, StringArray, StringDictionaryBuilder, + TimestampMillisecondArray, UInt8Array, UInt64Array, + }; use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use futures::FutureExt; @@ -1654,7 +1517,7 @@ mod tests { assert_eq!(vec![0, 3, 2, 1], drained); } - fn drain_tournament_tree(tree: &mut TournamentTree) -> Vec { + fn drain_winner_tree(tree: &mut WinnerTree) -> Vec { let mut values = Vec::with_capacity(tree.len()); while let Some(value) = tree.pop() { values.push(value); @@ -1663,8 +1526,8 @@ mod tests { } #[test] - fn test_tournament_tree_empty() { - let mut tree = TournamentTree::::with_capacity(0); + fn test_winner_tree_empty() { + let mut tree = WinnerTree::::with_capacity(0); assert!(tree.is_empty()); assert_eq!(0, tree.len()); @@ -1676,8 +1539,8 @@ mod tests { } #[test] - fn test_tournament_tree_single_element() { - let mut tree = TournamentTree::with_capacity(1); + fn test_winner_tree_single_element() { + let mut tree = WinnerTree::with_capacity(1); tree.push(7); assert!(!tree.is_empty()); @@ -1689,35 +1552,32 @@ mod tests { } #[test] - fn test_tournament_tree_drains_in_descending_order() { - let mut tree = TournamentTree::with_capacity(8); + fn test_winner_tree_drains_in_descending_order() { + let mut tree = WinnerTree::with_capacity(8); for value in [3, 1, 4, 1, 5, 9, 2, 6] { tree.push(value); } assert_eq!(Some(&9), tree.peek()); assert_eq!(Some(&6), tree.second_best()); - assert_eq!( - vec![9, 6, 5, 4, 3, 2, 1, 1], - drain_tournament_tree(&mut tree) - ); + assert_eq!(vec![9, 6, 5, 4, 3, 2, 1, 1], drain_winner_tree(&mut tree)); } #[test] - fn test_tournament_tree_non_power_of_two_capacity() { - let mut tree = TournamentTree::with_capacity(5); + fn test_winner_tree_non_power_of_two_capacity() { + let mut tree = WinnerTree::with_capacity(5); for value in [40, 10, 50, 20, 30] { tree.push(value); } assert_eq!(Some(&50), tree.peek()); assert_eq!(Some(&40), tree.second_best()); - assert_eq!(vec![50, 40, 30, 20, 10], drain_tournament_tree(&mut tree)); + assert_eq!(vec![50, 40, 30, 20, 10], drain_winner_tree(&mut tree)); } #[test] - fn test_tournament_tree_replays_winner_after_mutation() { - let mut tree = TournamentTree::with_capacity(4); + fn test_winner_tree_replays_winner_after_mutation() { + let mut tree = WinnerTree::with_capacity(4); for value in [7, 3, 9, 5] { tree.push(value); } @@ -1727,12 +1587,12 @@ mod tests { assert_eq!(Some(&7), tree.peek()); assert_eq!(Some(&5), tree.second_best()); - assert_eq!(vec![7, 5, 3, 1], drain_tournament_tree(&mut tree)); + assert_eq!(vec![7, 5, 3, 1], drain_winner_tree(&mut tree)); } #[test] - fn test_tournament_tree_mutated_winner_can_stay_winner() { - let mut tree = TournamentTree::with_capacity(3); + fn test_winner_tree_mutated_winner_can_stay_winner() { + let mut tree = WinnerTree::with_capacity(3); for value in [1, 2, 9] { tree.push(value); } @@ -1741,12 +1601,12 @@ mod tests { tree.replay_winner(); assert_eq!(Some(&8), tree.peek()); - assert_eq!(vec![8, 2, 1], drain_tournament_tree(&mut tree)); + assert_eq!(vec![8, 2, 1], drain_winner_tree(&mut tree)); } #[test] - fn test_tournament_tree_remove_and_reinsert() { - let mut tree = TournamentTree::with_capacity(3); + fn test_winner_tree_remove_and_reinsert() { + let mut tree = WinnerTree::with_capacity(3); for value in [5, 9, 7] { tree.push(value); } @@ -1755,20 +1615,20 @@ mod tests { assert_eq!(Some(9), tree.pop()); tree.push(8); assert_eq!(Some(&8), tree.peek()); - assert_eq!(vec![8, 7, 5], drain_tournament_tree(&mut tree)); + assert_eq!(vec![8, 7, 5], drain_winner_tree(&mut tree)); // Refill after the tree was drained to empty. tree.push(4); tree.push(6); assert_eq!(Some(&6), tree.peek()); - assert_eq!(vec![6, 4], drain_tournament_tree(&mut tree)); + assert_eq!(vec![6, 4], drain_winner_tree(&mut tree)); } - /// Drives a TournamentTree and a std BinaryHeap oracle with the same seeded op + /// Drives a WinnerTree and a std BinaryHeap oracle with the same seeded op /// sequence (push / pop winner / mutate winner + replay) and compares /// observable behavior after every op. The number of live elements never /// exceeds `capacity`, mirroring how MergeAlgo uses the tree. - fn assert_tournament_tree_matches_oracle( + fn assert_winner_tree_matches_oracle( seed: u64, value_range: u32, capacity: usize, @@ -1778,7 +1638,7 @@ mod tests { use rand::{Rng, SeedableRng}; let mut rng = StdRng::seed_from_u64(seed); - let mut tree = TournamentTree::::with_capacity(capacity); + let mut tree = WinnerTree::::with_capacity(capacity); let mut oracle = BinaryHeap::::new(); let mut next_value = 0_u32; @@ -1820,27 +1680,27 @@ mod tests { while let Some(value) = oracle.pop() { oracle_values.push(value); } - assert_eq!(oracle_values, drain_tournament_tree(&mut tree)); + assert_eq!(oracle_values, drain_winner_tree(&mut tree)); } #[test] - fn test_tournament_tree_matches_binary_heap_oracle() { + fn test_winner_tree_matches_binary_heap_oracle() { for seed in [0x5eed, 0xdead_beef, 42] { - assert_tournament_tree_matches_oracle(seed, 1000, 13, 2000); + assert_winner_tree_matches_oracle(seed, 1000, 13, 2000); } } #[test] - fn test_tournament_tree_matches_oracle_with_duplicate_heavy_values() { + fn test_winner_tree_matches_oracle_with_duplicate_heavy_values() { // A tiny value range makes duplicates dominate, which exercises the // tie-breaking branches of the tree matches. - assert_tournament_tree_matches_oracle(0xc0ffee, 3, 8, 2000); + assert_winner_tree_matches_oracle(0xc0ffee, 3, 8, 2000); } #[test] - fn test_tournament_tree_matches_oracle_with_tiny_capacities() { + fn test_winner_tree_matches_oracle_with_tiny_capacities() { for capacity in 1..=3 { - assert_tournament_tree_matches_oracle(0xbeef, 100, capacity, 500); + assert_winner_tree_matches_oracle(0xbeef, 100, capacity, 500); } } @@ -1937,7 +1797,7 @@ mod tests { &[OpType::Put, OpType::Put], &[11, 12], ); - let mut cursor = RowCursor::new(SortColumns::new(&batch)); + let mut cursor = RowCursor::new(SortColumns::try_new(&batch).unwrap()); assert!(!cursor.is_last_row()); cursor.advance(); @@ -2114,6 +1974,167 @@ mod tests { assert_record_batches_eq(&expected, &result); } + #[test] + fn test_merge_iterator_empty_primary_keys() { + // Tables without tags produce batches whose primary keys are all empty + // byte strings. + let batch1 = create_test_record_batch( + &[b"", b""], + &[1000, 3000], + &[21, 23], + &[OpType::Put, OpType::Put], + &[11, 13], + ); + let batch2 = create_test_record_batch( + &[b"", b""], + &[2000, 4000], + &[22, 24], + &[OpType::Put, OpType::Put], + &[12, 14], + ); + + let schema = batch1.schema(); + let iter1 = Box::new(new_test_iter(vec![batch1])); + let iter2 = Box::new(new_test_iter(vec![batch2])); + + let merge_iter = FlatMergeIterator::new(schema, vec![iter1, iter2], 1024).unwrap(); + let result = collect_merge_iterator_batches(merge_iter); + + let num_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(4, num_rows); + let mut timestamps = Vec::new(); + for batch in &result { + let pk_idx = primary_key_column_index(batch.num_columns()); + // All rows share the same empty primary key, so each output + // dictionary must contain a single empty value. + let expected_keys = vec![b"".as_slice(); batch.num_rows()]; + assert_primary_key_dictionary(batch.column(pk_idx).as_ref(), &expected_keys, &[b""]); + let timestamp = batch + .column(time_index_column_index(batch.num_columns())) + .as_any() + .downcast_ref::() + .unwrap(); + timestamps.extend(timestamp.values().iter().copied()); + } + // Rows are merged by timestamp since all primary keys are equal. + assert_eq!(vec![1000, 2000, 3000, 4000], timestamps); + } + + /// Creates a test RecordBatch with an extra dictionary-encoded string tag + /// column, mirroring the flat input schema of tables with string tags. + fn create_test_record_batch_with_dict_tag( + tags: &[&str], + primary_keys: &[&[u8]], + timestamps: &[i64], + sequences: &[u64], + op_types: &[OpType], + field_values: &[i64], + ) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "tag0", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + Field::new("field1", DataType::Int64, false), + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + "__primary_key", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)), + false, + ), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])); + + let mut tag_builder = StringDictionaryBuilder::::new(); + for tag in tags { + tag_builder.append(tag).unwrap(); + } + let tag = Arc::new(tag_builder.finish()); + + let field1 = Arc::new(Int64Array::from_iter_values(field_values.iter().copied())); + let timestamp = Arc::new(TimestampMillisecondArray::from_iter_values( + timestamps.iter().copied(), + )); + let mut pk_builder = BinaryDictionaryBuilder::::new(); + for key in primary_keys { + pk_builder.append(key).unwrap(); + } + let primary_key = Arc::new(pk_builder.finish()); + let sequence = Arc::new(UInt64Array::from_iter_values(sequences.iter().copied())); + let op_type = Arc::new(UInt8Array::from_iter_values( + op_types.iter().map(|&v| v as u8), + )); + + RecordBatch::try_new( + schema, + vec![tag, field1, timestamp, primary_key, sequence, op_type], + ) + .unwrap() + } + + #[test] + fn test_merge_iterator_dictionary_encoded_tag_column() { + // String tag columns are dictionary-encoded in the flat input schema, + // and each source may carry a different dictionary. + let batch1 = create_test_record_batch_with_dict_tag( + &["us-west", "eu-west"], + &[b"k1", b"k2"], + &[1000, 2000], + &[21, 22], + &[OpType::Put, OpType::Put], + &[11, 12], + ); + let batch2 = create_test_record_batch_with_dict_tag( + &["us-east", "eu-west"], + &[b"k1", b"k2"], + &[1500, 2500], + &[23, 24], + &[OpType::Put, OpType::Put], + &[13, 14], + ); + + let schema = batch1.schema(); + let iter1 = Box::new(new_test_iter(vec![batch1])); + let iter2 = Box::new(new_test_iter(vec![batch2])); + + let merge_iter = FlatMergeIterator::new(schema, vec![iter1, iter2], 1024).unwrap(); + let result = collect_merge_iterator_batches(merge_iter); + + // Rows merged by (primary key, timestamp): (k1, 1000), (k1, 1500), + // (k2, 2000), (k2, 2500). + let num_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(4, num_rows); + let mut tags = Vec::new(); + let mut timestamps = Vec::new(); + for batch in &result { + let tag = batch + .column(0) + .as_any() + .downcast_ref::>() + .unwrap(); + let tag_values = tag.values().as_any().downcast_ref::().unwrap(); + tags.extend( + tag.keys() + .iter() + .map(|key| tag_values.value(key.unwrap() as usize)), + ); + let timestamp = batch + .column(time_index_column_index(batch.num_columns())) + .as_any() + .downcast_ref::() + .unwrap(); + timestamps.extend(timestamp.values().iter().copied()); + } + assert_eq!(vec!["us-west", "us-east", "eu-west", "eu-west"], tags); + assert_eq!(vec![1000, 1500, 2000, 2500], timestamps); + } + #[test] fn test_merge_iterator_retry_after_row_boundary_error_removes_source() { let (first, second, pending) = boundary_test_batches(); @@ -2244,6 +2265,33 @@ mod tests { assert_eq!(batch, output_batch); } + #[test] + fn test_merge_iterator_rejects_batch_without_internal_columns() { + // A generic schema without the flat-format internal columns cannot + // drive row comparison; the merger must return an error instead of + // panicking. + let schema = Arc::new(Schema::new(vec![ + Field::new("field1", DataType::Int64, false), + Field::new("field2", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(Int64Array::from(vec![3, 4])), + ], + ) + .unwrap(); + let iter = Box::new(new_test_iter(vec![batch])); + + let result = FlatMergeIterator::new(schema, vec![iter], 1024); + + assert!(matches!( + result, + Err(crate::error::Error::InvalidRecordBatch { .. }) + )); + } + fn assert_primary_key_dictionary( array: &dyn Array, expected_decoded: &[&[u8]], @@ -2367,8 +2415,8 @@ mod tests { &[11, 12], ); - let columns1 = SortColumns::new(&batch1); - let columns2 = SortColumns::new(&batch2); + let columns1 = SortColumns::try_new(&batch1).unwrap(); + let columns2 = SortColumns::try_new(&batch2).unwrap(); let cursor1 = RowCursor::new(columns1); let cursor2 = RowCursor::new(columns2); @@ -2382,8 +2430,8 @@ mod tests { fn test_row_cursor_caches_current_primary_key() { let batch1 = create_test_record_batch(&[b"k1"], &[1000], &[1], &[OpType::Put], &[11]); let batch2 = create_test_record_batch(&[b"k2"], &[1000], &[1], &[OpType::Put], &[12]); - let cursor1 = RowCursor::new(SortColumns::new(&batch1)); - let cursor2 = RowCursor::new(SortColumns::new(&batch2)); + let cursor1 = RowCursor::new(SortColumns::try_new(&batch1).unwrap()); + let cursor2 = RowCursor::new(SortColumns::try_new(&batch2).unwrap()); for _ in 0..5 { assert_eq!(Ordering::Less, cursor1.cmp(&cursor2));