rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-10 07:22:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mito2): harden flat merge and use winner_tree (#9064)

Browse Source 
* fix(mito2): reject non-flat batches in flat merge instead of panicking

- SortColumns::new becomes fallible try_new: batches missing the
  flat-format internal columns (time index, __primary_key, __sequence)
  at the fixed trailing positions now yield InvalidRecordBatch instead
  of a downcast panic, completing the generic-schema gate that only
  covered BatchBuilder output assembly. Document the flat-format input
  contract on FlatMergeIterator/FlatMergeReader.
- Clarify why BatchBuilder's schema gate uses >= 3 columns when a real
  flat-format schema always has at least 4.
- Add schema-structure tests: empty primary keys (tables without tags),
  dictionary-encoded string tag columns with per-source dictionaries,
  and graceful rejection of batches without internal columns.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* test(mito2): scan tables with various schemas through flat merge

Add an engine-level test that writes, flushes and scans regions without
tags (empty primary key) and with multiple string tags (dictionary-encoded
in the flat input schema), so the flat merge reader merges an SST with
the memtable on real schemas instead of hand-built batches.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* refactor(mito2): use winner_tree dependency

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* Update comments for FlatMergeIterator struct

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

---------

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>
Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Lei, HUANG
2026-09-08 07:20:18 +00:00
committed by GitHub
co-authored by Copilot Autofix powered by AI
parent 42a58080ea
commit be482553b0
5 changed files with 394 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+7
View File
@@ -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"
Cargo.toml
+1
View File
@@ -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
src/mito2/Cargo.toml
+1
View File
@@ -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]
src/mito2/src/engine/basic_test.rs
+86
View File
@@ -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<api::v1::Row> {
rows.iter()
.map(|&(value, ts)| {
let mut values: Vec<ValueData> = (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;
src/mito2/src/read/flat_merge.rs
+299 -251
View File
@@ -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<T> {
/// Slot storage, one slot per node. `None` means the slot is empty.
nodes: Vec<Option<T>>,
/// 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<usize>,
/// 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<T: Ord> TournamentTree<T> {
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<T> {
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<usize> {
(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<usize> {
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<usize> {
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<T: Ord> {
///
/// 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<T>,
hot: WinnerTree<T>,
/// 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<T: NodeCmp> MergeAlgo<T> {
fn new(mut nodes: Vec<T>) -> 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<Self> {
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::<PrimaryKeyArray>()
.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::<BinaryArray>()
.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::<UInt64Array>()
.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<IterNode>,
@@ -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<StreamNode>,
@@ -1316,7 +1176,7 @@ impl GenericNode<BoxedRecordBatchIterator> {
/// Returns the fetched new batch.
fn advance_batch(&mut self) -> Result<Option<RecordBatch>> {
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<BoxedRecordBatchStream> {
/// Returns the fetched new batch.
async fn advance_batch(&mut self) -> Result<Option<RecordBatch>> {
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<T: Ord>(tree: &mut TournamentTree<T>) -> Vec<T> {
fn drain_winner_tree<T: Ord>(tree: &mut WinnerTree<T>) -> Vec<T> {
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::<i32>::with_capacity(0);
fn test_winner_tree_empty() {
let mut tree = WinnerTree::<i32>::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::<u32>::with_capacity(capacity);
let mut tree = WinnerTree::<u32>::with_capacity(capacity);
let mut oracle = BinaryHeap::<u32>::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::<TimestampMillisecondArray>()
.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::<UInt32Type>::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::<UInt32Type>::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::<DictionaryArray<UInt32Type>>()
.unwrap();
let tag_values = tag.values().as_any().downcast_ref::<StringArray>().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::<TimestampMillisecondArray>()
.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));