rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-30 15:36:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: change the receivers of merge tree components (#3378)

Browse Source 
* refactor: change the receivers of Shard::read/DataBuffer::read/DataParts::read to &self instead of &mut self

* refactor: remove allow(dead_code) in merge tree
This commit is contained in:
Lei, HUANG
2024-02-26 06:50:55 +00:00
committed by GitHub
parent e481f073f5
commit 8087822ab2
8 changed files with 131 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/memtable.rs
-1
View File
@@ -15,7 +15,6 @@
//! Memtables are write buffers for regions.
pub mod key_values;
#[allow(dead_code)]
pub mod merge_tree;
pub mod time_series;
pub(crate) mod version;
src/mito2/src/memtable/merge_tree/data.rs
+119 -106
View File
@@ -30,7 +30,7 @@ use datatypes::types::TimestampType;
use datatypes::vectors::{
TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector,
TimestampSecondVector, UInt16Vector, UInt16VectorBuilder, UInt64Vector, UInt64VectorBuilder,
UInt8VectorBuilder,
UInt8Vector, UInt8VectorBuilder,
};
use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
use parquet::arrow::ArrowWriter;
@@ -65,10 +65,6 @@ impl DataBatchRange {
pub(crate) fn len(&self) -> usize {
self.end - self.start
}
pub(crate) fn is_empty(&self) -> bool {
self.len() == 0
}
}
/// Data part batches returns by `DataParts::read`.
@@ -89,10 +85,6 @@ impl<'a> DataBatch<'a> {
self.range
}
pub(crate) fn is_empty(&self) -> bool {
self.range.is_empty()
}
pub(crate) fn slice_record_batch(&self) -> RecordBatch {
self.rb.slice(self.range.start, self.range.len())
}
@@ -258,12 +250,11 @@ impl DataBuffer {
/// Reads batches from data buffer without resetting builder's buffers.
/// If pk_weights is present, yielded rows are sorted according to weights,
/// otherwise rows are sorted by "pk_weights" values as they are actually weights.
pub fn read(&mut self, pk_weights: Option<&[u16]>) -> Result<DataBufferReader> {
let batch = data_buffer_to_record_batches(
pub fn read(&self, pk_weights: Option<&[u16]>) -> Result<DataBufferReader> {
let batch = read_data_buffer_to_record_batches(
self.data_part_schema.clone(),
self,
pk_weights,
true,
self.dedup,
// replace_pk_index is always set to false since:
// - for DataBuffer in ShardBuilder, pk dict is not frozen
@@ -307,36 +298,115 @@ impl LazyMutableVectorBuilder {
}
/// Converts `DataBuffer` to record batches, with rows sorted according to pk_weights.
/// `keep_data`: whether to keep the original data inside `DataBuffer`.
/// `dedup`: whether to true to remove the duplicated rows inside `DataBuffer`.
/// `replace_pk_index`: whether to replace the pk_index values with corresponding pk weight.
fn data_buffer_to_record_batches(
fn drain_data_buffer_to_record_batches(
schema: SchemaRef,
buffer: &mut DataBuffer,
pk_weights: Option<&[u16]>,
keep_data: bool,
dedup: bool,
replace_pk_index: bool,
) -> Result<RecordBatch> {
let num_rows = buffer.ts_builder.len();
let (pk_index_v, ts_v, sequence_v, op_type_v) = if keep_data {
(
buffer.pk_index_builder.finish_cloned(),
buffer.ts_builder.to_vector_cloned(),
buffer.sequence_builder.finish_cloned(),
buffer.op_type_builder.finish_cloned(),
)
} else {
(
buffer.pk_index_builder.finish(),
buffer.ts_builder.to_vector(),
buffer.sequence_builder.finish(),
buffer.op_type_builder.finish(),
)
};
let (pk_index_v, ts_v, sequence_v, op_type_v) = (
buffer.pk_index_builder.finish(),
buffer.ts_builder.to_vector(),
buffer.sequence_builder.finish(),
buffer.op_type_builder.finish(),
);
let mut rows = build_rows_to_sort(pk_weights, &pk_index_v, &ts_v, &sequence_v);
let (indices_to_take, mut columns) = build_row_sort_indices_and_columns(
pk_weights,
pk_index_v,
ts_v,
sequence_v,
op_type_v,
replace_pk_index,
dedup,
buffer.field_builders.len() + 4,
)?;
for b in buffer.field_builders.iter_mut() {
let array = match b {
LazyMutableVectorBuilder::Type(ty) => {
let mut single_null = ty.create_mutable_vector(num_rows);
single_null.push_nulls(num_rows);
single_null.to_vector().to_arrow_array()
}
LazyMutableVectorBuilder::Builder(builder) => builder.to_vector().to_arrow_array(),
};
columns.push(
arrow::compute::take(&array, &indices_to_take, None)
.context(error::ComputeArrowSnafu)?,
);
}
RecordBatch::try_new(schema, columns).context(error::NewRecordBatchSnafu)
}
/// Reads `DataBuffer` to record batches, with rows sorted according to pk_weights without resetting `DataBuffer`.
/// `dedup`: whether to true to remove the duplicated rows inside `DataBuffer`.
/// `replace_pk_index`: whether to replace the pk_index values with corresponding pk weight.
fn read_data_buffer_to_record_batches(
schema: SchemaRef,
buffer: &DataBuffer,
pk_weights: Option<&[u16]>,
dedup: bool,
replace_pk_index: bool,
) -> Result<RecordBatch> {
let num_rows = buffer.ts_builder.len();
let (pk_index_v, ts_v, sequence_v, op_type_v) = (
buffer.pk_index_builder.finish_cloned(),
buffer.ts_builder.to_vector_cloned(),
buffer.sequence_builder.finish_cloned(),
buffer.op_type_builder.finish_cloned(),
);
let (indices_to_take, mut columns) = build_row_sort_indices_and_columns(
pk_weights,
pk_index_v,
ts_v,
sequence_v,
op_type_v,
replace_pk_index,
dedup,
buffer.field_builders.len() + 4,
)?;
for b in buffer.field_builders.iter() {
let array = match b {
LazyMutableVectorBuilder::Type(ty) => {
let mut single_null = ty.create_mutable_vector(num_rows);
single_null.push_nulls(num_rows);
single_null.to_vector().to_arrow_array()
}
LazyMutableVectorBuilder::Builder(builder) => {
builder.to_vector_cloned().to_arrow_array()
}
};
columns.push(
arrow::compute::take(&array, &indices_to_take, None)
.context(error::ComputeArrowSnafu)?,
);
}
RecordBatch::try_new(schema, columns).context(error::NewRecordBatchSnafu)
}
#[allow(clippy::too_many_arguments)]
fn build_row_sort_indices_and_columns(
pk_weights: Option<&[u16]>,
pk_index: UInt16Vector,
ts: VectorRef,
sequence: UInt64Vector,
op_type: UInt8Vector,
replace_pk_index: bool,
dedup: bool,
column_num: usize,
) -> Result<(UInt32Array, Vec<ArrayRef>)> {
let mut rows = build_rows_to_sort(pk_weights, &pk_index, &ts, &sequence);
let pk_array = if replace_pk_index {
// replace pk index values with pk weights.
@@ -344,7 +414,7 @@ fn data_buffer_to_record_batches(
rows.iter().map(|(_, key)| key.pk_weight),
)) as Arc<_>
} else {
pk_index_v.to_arrow_array()
pk_index.to_arrow_array()
};
// sort and dedup
@@ -355,7 +425,7 @@ fn data_buffer_to_record_batches(
let indices_to_take = UInt32Array::from_iter_values(rows.iter().map(|(idx, _)| *idx as u32));
let mut columns = Vec::with_capacity(4 + buffer.field_builders.len());
let mut columns = Vec::with_capacity(column_num);
columns.push(
arrow::compute::take(&pk_array, &indices_to_take, None)
@@ -363,42 +433,21 @@ fn data_buffer_to_record_batches(
);
columns.push(
arrow::compute::take(&ts_v.to_arrow_array(), &indices_to_take, None)
arrow::compute::take(&ts.to_arrow_array(), &indices_to_take, None)
.context(error::ComputeArrowSnafu)?,
);
columns.push(
arrow::compute::take(&sequence_v.as_arrow(), &indices_to_take, None)
arrow::compute::take(&sequence.as_arrow(), &indices_to_take, None)
.context(error::ComputeArrowSnafu)?,
);
columns.push(
arrow::compute::take(&op_type_v.as_arrow(), &indices_to_take, None)
arrow::compute::take(&op_type.as_arrow(), &indices_to_take, None)
.context(error::ComputeArrowSnafu)?,
);
for b in buffer.field_builders.iter_mut() {
let array = match b {
LazyMutableVectorBuilder::Type(ty) => {
let mut single_null = ty.create_mutable_vector(num_rows);
single_null.push_nulls(num_rows);
single_null.to_vector().to_arrow_array()
}
LazyMutableVectorBuilder::Builder(builder) => {
if keep_data {
builder.to_vector_cloned().to_arrow_array()
} else {
builder.to_vector().to_arrow_array()
}
}
};
columns.push(
arrow::compute::take(&array, &indices_to_take, None)
.context(error::ComputeArrowSnafu)?,
);
}
RecordBatch::try_new(schema, columns).context(error::NewRecordBatchSnafu)
Ok((indices_to_take, columns))
}
pub(crate) fn timestamp_array_to_i64_slice(arr: &ArrayRef) -> &[i64] {
@@ -468,12 +517,6 @@ impl DataBufferReader {
}
}
/// # Panics
/// If Current reader is exhausted.
pub(crate) fn current_pk_index(&self) -> PkIndex {
self.current_range.as_ref().unwrap().pk_index
}
/// Advances reader to next data batch.
pub(crate) fn next(&mut self) -> Result<()> {
if self.offset >= self.batch.num_rows() {
@@ -673,11 +716,10 @@ impl<'a> DataPartEncoder<'a> {
let mut bytes = Vec::with_capacity(1024);
let mut writer = ArrowWriter::try_new(&mut bytes, self.schema.clone(), self.writer_props())
.context(error::EncodeMemtableSnafu)?;
let rb = data_buffer_to_record_batches(
let rb = drain_data_buffer_to_record_batches(
self.schema.clone(),
source,
self.pk_weights,
false,
self.dedup,
self.replace_pk_index,
)?;
@@ -695,12 +737,6 @@ pub enum DataPart {
}
impl DataPart {
fn is_empty(&self) -> bool {
match self {
DataPart::Parquet(p) => p.data.is_empty(),
}
}
/// Reads frozen data part and yields [DataBatch]es.
pub fn read(&self) -> Result<DataPartReader> {
match self {
@@ -745,14 +781,6 @@ impl DataPartReader {
self.current_range.is_some()
}
/// Returns current pk index.
///
/// # Panics
/// If reader is exhausted.
pub(crate) fn current_pk_index(&self) -> PkIndex {
self.current_range.as_ref().unwrap().pk_index
}
/// Returns current data batch of reader.
/// # Panics
/// If reader is exhausted.
@@ -835,16 +863,10 @@ impl DataParts {
self.active.write_row(pk_index, kv)
}
/// Freezes the active data buffer into frozen data parts.
pub fn freeze(&mut self) -> Result<()> {
self.frozen.push(self.active.freeze(None, false)?);
Ok(())
}
/// Reads data from all parts including active and frozen parts.
/// The returned iterator yields a record batch of one primary key at a time.
/// The order of yielding primary keys is determined by provided weights.
pub fn read(&mut self) -> Result<DataPartsReader> {
pub fn read(&self) -> Result<DataPartsReader> {
let mut nodes = Vec::with_capacity(self.frozen.len() + 1);
nodes.push(DataNode::new(DataSource::Buffer(
// `DataPars::read` ensures that all pk_index inside `DataBuffer` are replaced by weights.
@@ -857,10 +879,6 @@ impl DataParts {
let merger = Merger::try_new(nodes)?;
Ok(DataPartsReader { merger })
}
pub(crate) fn is_empty(&self) -> bool {
self.active.is_empty() && self.frozen.iter().all(|part| part.is_empty())
}
}
/// Reader for all parts inside a `DataParts`.
@@ -924,15 +942,12 @@ mod tests {
write_rows_to_buffer(&mut buffer, &meta, 0, vec![2], vec![Some(1.1)], 3);
assert_eq!(5, buffer.num_rows());
let schema = memtable_schema_to_encoded_schema(&meta);
let batch = data_buffer_to_record_batches(
schema,
&mut buffer,
Some(&[3, 1]),
keep_data,
true,
true,
)
.unwrap();
let batch = if keep_data {
read_data_buffer_to_record_batches(schema, &buffer, Some(&[3, 1]), true, true).unwrap()
} else {
drain_data_buffer_to_record_batches(schema, &mut buffer, Some(&[3, 1]), true, true)
.unwrap()
};
assert_eq!(
vec![1, 2, 1, 2],
@@ -1036,8 +1051,7 @@ mod tests {
assert_eq!(4, buffer.num_rows());
let schema = memtable_schema_to_encoded_schema(&meta);
let batch =
data_buffer_to_record_batches(schema, &mut buffer, Some(&[0, 1]), true, true, true)
.unwrap();
read_data_buffer_to_record_batches(schema, &buffer, Some(&[0, 1]), true, true).unwrap();
assert_eq!(3, batch.num_rows());
assert_eq!(
@@ -1090,9 +1104,8 @@ mod tests {
write_rows_to_buffer(&mut buffer, &meta, 0, vec![2], vec![Some(1.1)], 3);
assert_eq!(5, buffer.num_rows());
let schema = memtable_schema_to_encoded_schema(&meta);
let batch =
data_buffer_to_record_batches(schema, &mut buffer, Some(&[3, 1]), true, false, true)
.unwrap();
let batch = read_data_buffer_to_record_batches(schema, &buffer, Some(&[3, 1]), false, true)
.unwrap();
assert_eq!(
vec![1, 1, 3, 3, 3],
@@ -1289,7 +1302,7 @@ mod tests {
#[test]
fn test_iter_empty_data_buffer() {
let meta = metadata_for_test();
let mut buffer = DataBuffer::with_capacity(meta.clone(), 10, true);
let buffer = DataBuffer::with_capacity(meta.clone(), 10, true);
let mut iter = buffer.read(Some(&[0, 1, 3, 2])).unwrap();
check_buffer_values_equal(&mut iter, &[]);
}
src/mito2/src/memtable/merge_tree/dedup.rs
+1 -1
View File
@@ -176,7 +176,7 @@ mod tests {
frozens.push(part1);
}
let mut parts = DataParts::new(meta, 10, true).with_frozen(frozens);
let parts = DataParts::new(meta, 10, true).with_frozen(frozens);
let mut res = Vec::with_capacity(expected.len());
let mut reader = DedupReader::try_new(MockSource(parts.read().unwrap())).unwrap();
src/mito2/src/memtable/merge_tree/dict.rs
+6 -33
View File
@@ -56,11 +56,6 @@ impl KeyDictBuilder {
}
}
/// Gets the pk index by the key.
pub fn get_pk_index(&self, key: &[u8]) -> Option<PkIndex> {
self.pk_to_index.get(key).copied()
}
/// Returns true if the builder is full.
pub fn is_full(&self) -> bool {
self.num_keys >= self.capacity
@@ -97,6 +92,7 @@ impl KeyDictBuilder {
}
/// Memory size of the builder.
#[cfg(test)]
pub fn memory_size(&self) -> usize {
self.key_bytes_in_index
+ self.key_buffer.buffer_memory_size()
@@ -152,8 +148,6 @@ impl KeyDictBuilder {
pub struct DictBuilderReader {
blocks: Vec<DictBlock>,
sorted_pk_indices: Vec<PkIndex>,
/// Current offset in the `sorted_pk_indices`.
offset: usize,
}
impl DictBuilderReader {
@@ -161,21 +155,23 @@ impl DictBuilderReader {
Self {
blocks,
sorted_pk_indices,
offset: 0,
}
}
/// Returns the number of keys.
#[cfg(test)]
pub fn num_keys(&self) -> usize {
self.sorted_pk_indices.len()
}
/// Gets the i-th pk index.
#[cfg(test)]
pub fn pk_index(&self, offset: usize) -> PkIndex {
self.sorted_pk_indices[offset]
}
/// Gets the i-th key.
#[cfg(test)]
pub fn key(&self, offset: usize) -> &[u8] {
let pk_index = self.pk_index(offset);
self.key_by_pk_index(pk_index)
@@ -191,11 +187,6 @@ impl DictBuilderReader {
pub(crate) fn pk_weights_to_sort_data(&self, pk_weights: &mut Vec<u16>) {
compute_pk_weights(&self.sorted_pk_indices, pk_weights)
}
/// Returns pk indices sorted by keys.
pub(crate) fn sorted_pk_index(&self) -> &[PkIndex] {
&self.sorted_pk_indices
}
}
/// Returns pk weights to sort a data part and replaces pk indices.
@@ -290,23 +281,8 @@ impl KeyBuffer {
self.key_builder.is_empty()
}
/// Gets the primary key by its index.
///
/// # Panics
/// Panics if the index is invalid.
fn get_key(&self, index: PkIndex) -> &[u8] {
let values = self.key_builder.values_slice();
let offsets = self.key_builder.offsets_slice();
// Casting index to usize is safe.
let start = offsets[index as usize];
let end = offsets[index as usize + 1];
// We ensure no null in the builder so we don't check validity.
// The builder offset should be positive.
&values[start as usize..end as usize]
}
/// Returns the buffer size of the builder.
#[cfg(test)]
fn buffer_memory_size(&self) -> usize {
self.key_builder.values_slice().len()
+ std::mem::size_of_val(self.key_builder.offsets_slice())
@@ -351,15 +327,12 @@ impl DictBlock {
Self { keys }
}
fn len(&self) -> usize {
self.keys.len()
}
fn key_by_pk_index(&self, index: PkIndex) -> &[u8] {
let pos = index % MAX_KEYS_PER_BLOCK;
self.keys.value(pos as usize)
}
#[cfg(test)]
fn buffer_memory_size(&self) -> usize {
self.keys.get_buffer_memory_size()
}
src/mito2/src/memtable/merge_tree/partition.rs
+2 -3
View File
@@ -105,13 +105,12 @@ impl Partition {
/// Scans data in the partition.
pub fn read(&self, mut context: ReadPartitionContext) -> Result<PartitionReader> {
// TODO(yingwen): Change to acquire read lock if `read()` takes `&self`.
let nodes = {
let mut inner = self.inner.write().unwrap();
let inner = self.inner.read().unwrap();
let mut nodes = Vec::with_capacity(inner.shards.len() + 1);
let bulder_reader = inner.shard_builder.read(&mut context.pk_weights)?;
nodes.push(ShardNode::new(ShardSource::Builder(bulder_reader)));
for shard in &mut inner.shards {
for shard in &inner.shards {
let shard_reader = shard.read()?;
nodes.push(ShardNode::new(ShardSource::Shard(shard_reader)));
}
src/mito2/src/memtable/merge_tree/shard.rs
+1 -5
View File
@@ -72,7 +72,7 @@ impl Shard {
/// Scans the shard.
// TODO(yingwen): Push down projection to data parts.
pub fn read(&mut self) -> Result<ShardReader> {
pub fn read(&self) -> Result<ShardReader> {
let parts_reader = self.data_parts.read()?;
Ok(ShardReader {
@@ -136,10 +136,6 @@ pub struct ShardReader {
}
impl ShardReader {
fn shard_id(&self) -> ShardId {
self.shard_id
}
fn is_valid(&self) -> bool {
self.parts_reader.is_valid()
}
src/mito2/src/memtable/merge_tree/shard_builder.rs
+2 -29
View File
@@ -106,7 +106,7 @@ impl ShardBuilder {
}
/// Scans the shard builder.
pub fn read(&mut self, pk_weights_buffer: &mut Vec<u16>) -> Result<ShardBuilderReader> {
pub fn read(&self, pk_weights_buffer: &mut Vec<u16>) -> Result<ShardBuilderReader> {
let dict_reader = self.dict_builder.read();
dict_reader.pk_weights_to_sort_data(pk_weights_buffer);
let data_reader = self.data_buffer.read(Some(pk_weights_buffer))?;
@@ -132,10 +132,6 @@ pub struct ShardBuilderReader {
}
impl ShardBuilderReader {
pub fn shard_id(&self) -> ShardId {
self.shard_id
}
pub fn is_valid(&self) -> bool {
self.data_reader.is_valid()
}
@@ -164,14 +160,12 @@ impl ShardBuilderReader {
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::*;
use crate::memtable::merge_tree::dict::KeyDictBuilder;
use crate::memtable::merge_tree::metrics::WriteMetrics;
use crate::memtable::KeyValues;
use crate::test_util::memtable_util::{
build_key_values_with_ts_seq_values, encode_key_by_kv, encode_keys, metadata_for_test,
build_key_values_with_ts_seq_values, encode_key_by_kv, metadata_for_test,
};
fn input_with_key(metadata: &RegionMetadataRef) -> Vec<KeyValues> {
@@ -203,27 +197,6 @@ mod tests {
]
}
fn new_shard_builder(
shard_id: ShardId,
metadata: RegionMetadataRef,
input: &[KeyValues],
) -> Shard {
let mut dict_builder = KeyDictBuilder::new(1024);
let mut metrics = WriteMetrics::default();
let mut keys = Vec::with_capacity(input.len());
for kvs in input {
encode_keys(&metadata, kvs, &mut keys);
}
for key in &keys {
dict_builder.insert_key(key, &mut metrics);
}
let dict = dict_builder.finish().unwrap();
let data_parts = DataParts::new(metadata, DATA_INIT_CAP, true);
Shard::new(shard_id, Some(Arc::new(dict)), data_parts, true)
}
#[test]
fn test_write_shard_builder() {
let metadata = metadata_for_test();
src/mito2/src/memtable/merge_tree/tree.rs
-17
View File
@@ -21,8 +21,6 @@ use api::v1::OpType;
use common_recordbatch::filter::SimpleFilterEvaluator;
use common_time::Timestamp;
use datafusion_common::ScalarValue;
use datatypes::arrow;
use datatypes::data_type::ConcreteDataType;
use snafu::ensure;
use store_api::metadata::RegionMetadataRef;
use store_api::storage::ColumnId;
@@ -35,7 +33,6 @@ use crate::memtable::merge_tree::partition::{
Partition, PartitionKey, PartitionReader, PartitionRef, ReadPartitionContext,
};
use crate::memtable::merge_tree::MergeTreeConfig;
use crate::memtable::time_series::primary_key_schema;
use crate::memtable::{BoxedBatchIterator, KeyValues};
use crate::read::Batch;
use crate::row_converter::{McmpRowCodec, RowCodec, SortField};
@@ -144,18 +141,8 @@ impl MergeTree {
.unwrap_or_default();
let partitions = self.prune_partitions(&filters);
let pk_schema = primary_key_schema(&self.metadata);
let pk_datatypes = self
.metadata
.primary_key_columns()
.map(|pk| pk.column_schema.data_type.clone())
.collect();
let mut iter = TreeIter {
metadata: self.metadata.clone(),
pk_schema,
pk_datatypes,
row_codec: self.row_codec.clone(),
partitions,
current_reader: None,
};
@@ -283,10 +270,6 @@ impl MergeTree {
}
struct TreeIter {
metadata: RegionMetadataRef,
pk_schema: arrow::datatypes::SchemaRef,
pk_datatypes: Vec<ConcreteDataType>,
row_codec: Arc<McmpRowCodec>,
partitions: VecDeque<PartitionRef>,
current_reader: Option<PartitionReader>,
}