mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-26 21:56:17 +00:00
refactor: isolate logical batch conversion and region writes
Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
WenyXu
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,686 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arrow::compute::concat_batches;
|
||||
use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use common_query::prelude::{greptime_timestamp, greptime_value};
|
||||
use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse};
|
||||
use smallvec::SmallVec;
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
use table::metadata::TableId;
|
||||
|
||||
use crate::error;
|
||||
use crate::error::Result;
|
||||
use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
|
||||
|
||||
pub(in crate::batcher::logical_table) const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3;
|
||||
|
||||
/// An aligned logical record batch and its timestamp column index.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RecordBatchWithTsIdx {
|
||||
/// The aligned logical record batch.
|
||||
pub(in crate::batcher::logical_table) batch: RecordBatch,
|
||||
/// The timestamp column index in `batch`.
|
||||
pub(in crate::batcher::logical_table) timestamp_index: usize,
|
||||
}
|
||||
|
||||
impl RecordBatchWithTsIdx {
|
||||
/// Creates a record batch with a validated timestamp column index.
|
||||
pub fn try_new(batch: RecordBatch, timestamp_index: usize) -> Result<Self> {
|
||||
let schema = batch.schema();
|
||||
let timestamp_field = schema.fields().get(timestamp_index).with_context(|| {
|
||||
error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Timestamp column index {} is out of bounds for record batch with {} columns",
|
||||
timestamp_index,
|
||||
batch.num_columns()
|
||||
),
|
||||
}
|
||||
})?;
|
||||
ensure!(
|
||||
matches!(timestamp_field.data_type(), ArrowDataType::Timestamp(_, _)),
|
||||
error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Column at index {} is not a timestamp column: {:?}",
|
||||
timestamp_index,
|
||||
timestamp_field.data_type()
|
||||
),
|
||||
}
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
batch,
|
||||
timestamp_index,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn into_parts(self) -> (RecordBatch, usize) {
|
||||
(self.batch, self.timestamp_index)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TableBatch {
|
||||
pub table_name: String,
|
||||
pub table_id: TableId,
|
||||
pub batches: Vec<RecordBatchWithTsIdx>,
|
||||
pub row_count: usize,
|
||||
}
|
||||
|
||||
/// Classifies columns in a logical-table batch for sparse primary-key conversion.
|
||||
///
|
||||
/// Returns:
|
||||
/// - `Vec<TagColumnInfo>`: all Utf8 tag columns sorted by tag name, used for
|
||||
/// TSID and sparse primary-key encoding.
|
||||
/// - `SmallVec<[usize; 3]>`: indices of columns copied into the physical batch
|
||||
/// after `__primary_key`, ordered as `[greptime_timestamp, greptime_value,
|
||||
/// partition_tag_columns...]`.
|
||||
pub(in crate::batcher::logical_table) fn columns_taxonomy(
|
||||
batch_schema: &Arc<ArrowSchema>,
|
||||
table_name: &str,
|
||||
name_to_ids: &HashMap<String, u32>,
|
||||
partition_columns: &HashSet<&str>,
|
||||
) -> Result<(Vec<TagColumnInfo>, SmallVec<[usize; 3]>)> {
|
||||
let mut tag_columns = Vec::new();
|
||||
let mut essential_column_indices =
|
||||
SmallVec::<[usize; 3]>::with_capacity(2 + partition_columns.len());
|
||||
// Placeholder for greptime_timestamp and greptime_value
|
||||
essential_column_indices.push(0);
|
||||
essential_column_indices.push(0);
|
||||
|
||||
let mut timestamp_index = None;
|
||||
let mut value_index = None;
|
||||
|
||||
for (index, field) in batch_schema.fields().iter().enumerate() {
|
||||
match field.data_type() {
|
||||
ArrowDataType::Utf8 => {
|
||||
let column_id = name_to_ids.get(field.name()).copied().with_context(|| {
|
||||
error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Column '{}' from logical table '{}' not found in physical table column IDs",
|
||||
field.name(),
|
||||
table_name
|
||||
),
|
||||
}
|
||||
})?;
|
||||
tag_columns.push(TagColumnInfo {
|
||||
name: field.name().clone(),
|
||||
index,
|
||||
column_id,
|
||||
});
|
||||
|
||||
if partition_columns.contains(field.name().as_str()) {
|
||||
essential_column_indices.push(index);
|
||||
}
|
||||
}
|
||||
ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => {
|
||||
ensure!(
|
||||
timestamp_index.replace(index).is_none(),
|
||||
error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Duplicated timestamp column in logical table '{}' batch schema",
|
||||
table_name
|
||||
),
|
||||
}
|
||||
);
|
||||
}
|
||||
ArrowDataType::Float64 => {
|
||||
ensure!(
|
||||
value_index.replace(index).is_none(),
|
||||
error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Duplicated value column in logical table '{}' batch schema",
|
||||
table_name
|
||||
),
|
||||
}
|
||||
);
|
||||
}
|
||||
datatype => {
|
||||
return error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Unexpected data type '{datatype:?}' in logical table '{}' batch schema",
|
||||
table_name
|
||||
),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let timestamp_index =
|
||||
timestamp_index.with_context(|| error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Missing essential column '{}' in logical table '{}' batch schema",
|
||||
greptime_timestamp(),
|
||||
table_name
|
||||
),
|
||||
})?;
|
||||
let value_index = value_index.with_context(|| error::InvalidPromRemoteRequestSnafu {
|
||||
msg: format!(
|
||||
"Missing essential column '{}' in logical table '{}' batch schema",
|
||||
greptime_value(),
|
||||
table_name
|
||||
),
|
||||
})?;
|
||||
|
||||
tag_columns.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
|
||||
essential_column_indices[0] = timestamp_index;
|
||||
essential_column_indices[1] = value_index;
|
||||
|
||||
Ok((tag_columns, essential_column_indices))
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn strip_partition_columns_from_batch(
|
||||
batch: RecordBatch,
|
||||
) -> Result<RecordBatch> {
|
||||
ensure!(
|
||||
batch.num_columns() >= PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT,
|
||||
error::InternalSnafu {
|
||||
err_msg: format!(
|
||||
"Expected at least {} columns in physical batch, got {}",
|
||||
PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT,
|
||||
batch.num_columns()
|
||||
),
|
||||
}
|
||||
);
|
||||
let essential_indices: Vec<usize> = (0..PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT).collect();
|
||||
batch.project(&essential_indices).context(error::ArrowSnafu)
|
||||
}
|
||||
|
||||
/// Transforms logical table batches into physical format (sparse primary key encoding).
|
||||
///
|
||||
/// It identifies tag columns and essential columns (timestamp, value) for each logical batch
|
||||
/// and applies sparse primary key modification.
|
||||
pub(in crate::batcher::logical_table) fn transform_logical_batches_to_physical(
|
||||
table_batches: &[TableBatch],
|
||||
name_to_ids: &HashMap<String, u32>,
|
||||
partition_columns_set: &HashSet<&str>,
|
||||
) -> Result<Vec<RecordBatch>> {
|
||||
let mut modified_batches: Vec<RecordBatch> =
|
||||
Vec::with_capacity(table_batches.iter().map(|b| b.batches.len()).sum());
|
||||
|
||||
let mut modify_elapsed = Duration::ZERO;
|
||||
let mut columns_taxonomy_elapsed = Duration::ZERO;
|
||||
|
||||
for table_batch in table_batches {
|
||||
let table_id = table_batch.table_id;
|
||||
|
||||
for batch in &table_batch.batches {
|
||||
let batch = &batch.batch;
|
||||
let batch_schema = batch.schema();
|
||||
let start = Instant::now();
|
||||
let (tag_columns, essential_col_indices) = columns_taxonomy(
|
||||
&batch_schema,
|
||||
&table_batch.table_name,
|
||||
name_to_ids,
|
||||
partition_columns_set,
|
||||
)?;
|
||||
|
||||
columns_taxonomy_elapsed += start.elapsed();
|
||||
if tag_columns.is_empty() && essential_col_indices.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let modified = {
|
||||
let start = Instant::now();
|
||||
// The schema of modified batch is: __primary_key, timestamp, value, other partition columns...
|
||||
let batch = modify_batch_sparse(
|
||||
batch.clone(),
|
||||
table_id,
|
||||
&tag_columns,
|
||||
&essential_col_indices,
|
||||
)?;
|
||||
modify_elapsed += start.elapsed();
|
||||
batch
|
||||
};
|
||||
|
||||
modified_batches.push(modified);
|
||||
}
|
||||
}
|
||||
|
||||
PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_modify_batch"])
|
||||
.observe(modify_elapsed.as_secs_f64());
|
||||
PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_columns_taxonomy"])
|
||||
.observe(columns_taxonomy_elapsed.as_secs_f64());
|
||||
|
||||
ensure!(
|
||||
!modified_batches.is_empty(),
|
||||
error::InternalSnafu {
|
||||
err_msg: "No batches can be transformed during pending flush",
|
||||
}
|
||||
);
|
||||
Ok(modified_batches)
|
||||
}
|
||||
|
||||
/// Concatenates all modified batches into a single large batch.
|
||||
///
|
||||
/// All modified batches share the same physical schema.
|
||||
pub(in crate::batcher::logical_table) fn concat_modified_batches(
|
||||
modified_batches: &[RecordBatch],
|
||||
) -> Result<RecordBatch> {
|
||||
let combined_schema = modified_batches[0].schema();
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_concat_all"])
|
||||
.start_timer();
|
||||
concat_batches(&combined_schema, modified_batches).context(error::ArrowSnafu)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow::array::{BinaryArray, StringArray, TimestampMillisecondArray};
|
||||
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use smallvec::SmallVec;
|
||||
|
||||
use crate::batcher::logical_table::batch_convert::{
|
||||
TableBatch, columns_taxonomy, strip_partition_columns_from_batch,
|
||||
transform_logical_batches_to_physical,
|
||||
};
|
||||
use crate::batcher::logical_table::test_util::mock_aligned_tag_batch;
|
||||
use crate::error::Error;
|
||||
|
||||
#[test]
|
||||
fn test_strip_partition_columns_from_batch_removes_partition_tags() {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("__primary_key", ArrowDataType::Binary, false),
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
|
||||
Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])),
|
||||
Arc::new(StringArray::from(vec!["node-1"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let stripped = strip_partition_columns_from_batch(batch).unwrap();
|
||||
|
||||
assert_eq!(3, stripped.num_columns());
|
||||
assert_eq!("__primary_key", stripped.schema().field(0).name());
|
||||
assert_eq!("greptime_timestamp", stripped.schema().field(1).name());
|
||||
assert_eq!("greptime_value", stripped.schema().field(2).name());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_partition_columns_from_batch_projects_essential_columns_without_lookup() {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("__primary_key", ArrowDataType::Binary, false),
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
|
||||
Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![42.0_f64])),
|
||||
Arc::new(StringArray::from(vec!["node-1"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let stripped = strip_partition_columns_from_batch(batch).unwrap();
|
||||
|
||||
assert_eq!(3, stripped.num_columns());
|
||||
assert_eq!("__primary_key", stripped.schema().field(0).name());
|
||||
assert_eq!("greptime_timestamp", stripped.schema().field(1).name());
|
||||
assert_eq!("greptime_value", stripped.schema().field(2).name());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_tag_columns_and_non_tag_indices_keeps_partition_tag_column() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
Field::new("region", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
let name_to_ids =
|
||||
HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]);
|
||||
let partition_columns = HashSet::from(["host"]);
|
||||
|
||||
let (tag_columns, non_tag_indices) =
|
||||
columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap();
|
||||
|
||||
assert_eq!(2, tag_columns.len());
|
||||
assert_eq!(&[0, 1, 2], non_tag_indices.as_slice());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_tag_columns_and_non_tag_indices_prioritizes_essential_columns() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("region", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
let name_to_ids =
|
||||
HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]);
|
||||
let partition_columns = HashSet::from(["host", "region"]);
|
||||
|
||||
let (_tag_columns, non_tag_indices): (_, SmallVec<[usize; 3]>) =
|
||||
columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap();
|
||||
|
||||
assert_eq!(&[2, 1, 0, 3], non_tag_indices.as_slice());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_tag_columns_and_non_tag_indices_rejects_unexpected_data_type() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
Field::new("invalid", ArrowDataType::Boolean, true),
|
||||
]));
|
||||
let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
|
||||
let partition_columns = HashSet::from(["host"]);
|
||||
|
||||
let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPromRemoteRequest { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_tag_columns_and_non_tag_indices_rejects_int64_timestamp_column() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("greptime_timestamp", ArrowDataType::Int64, false),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
|
||||
let partition_columns = HashSet::from(["host"]);
|
||||
|
||||
let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPromRemoteRequest { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_timestamp_column() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"ts1",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new(
|
||||
"ts2",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
|
||||
let partition_columns = HashSet::from(["host"]);
|
||||
|
||||
let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPromRemoteRequest { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_value_column() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("value1", ArrowDataType::Float64, true),
|
||||
Field::new("value2", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
|
||||
let partition_columns = HashSet::from(["host"]);
|
||||
|
||||
let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPromRemoteRequest { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_modify_batch_sparse_with_taxonomy_per_batch() {
|
||||
use arrow::array::BinaryArray;
|
||||
use metric_engine::batch_modifier::modify_batch_sparse;
|
||||
|
||||
let schema1 = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("tag1", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
|
||||
let schema2 = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("tag1", ArrowDataType::Utf8, true),
|
||||
Field::new("tag2", ArrowDataType::Utf8, true),
|
||||
]));
|
||||
let batch2 = RecordBatch::try_new(
|
||||
schema2.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondArray::from(vec![2000])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![2.0])),
|
||||
Arc::new(StringArray::from(vec!["v1"])),
|
||||
Arc::new(StringArray::from(vec!["v2"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]);
|
||||
let partition_columns = HashSet::new();
|
||||
|
||||
// A batch that only has tag1, same values as batch2 for ts and val.
|
||||
let batch3 = RecordBatch::try_new(
|
||||
schema1.clone(),
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondArray::from(vec![2000])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![2.0])),
|
||||
Arc::new(StringArray::from(vec!["v1"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Simulate the new loop logic in flush_batch_physical:
|
||||
// Resolve taxonomy FOR EACH BATCH.
|
||||
let (tag_columns2, indices2) =
|
||||
columns_taxonomy(&batch2.schema(), "table", &name_to_ids, &partition_columns).unwrap();
|
||||
let modified2 = modify_batch_sparse(batch2, 123, &tag_columns2, &indices2).unwrap();
|
||||
|
||||
let (tag_columns3, indices3) =
|
||||
columns_taxonomy(&batch3.schema(), "table", &name_to_ids, &partition_columns).unwrap();
|
||||
let modified3 = modify_batch_sparse(batch3, 123, &tag_columns3, &indices3).unwrap();
|
||||
|
||||
let pk2 = modified2
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap();
|
||||
let pk3 = modified3
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap();
|
||||
|
||||
// Now they SHOULD be different because tag2 is included in pk2 but not in pk3.
|
||||
assert_ne!(
|
||||
pk2.value(0),
|
||||
pk3.value(0),
|
||||
"PK should be different because batch2 has tag2!"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_logical_batches_to_physical_success() {
|
||||
let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
|
||||
|
||||
let table_batches = vec![TableBatch {
|
||||
table_name: "t1".to_string(),
|
||||
table_id: 1,
|
||||
batches: vec![batch],
|
||||
row_count: 1,
|
||||
}];
|
||||
|
||||
let name_to_ids = HashMap::from([("tag1".to_string(), 1)]);
|
||||
let partition_columns = HashSet::new();
|
||||
let modified =
|
||||
transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(1, modified.len());
|
||||
assert_eq!(3, modified[0].num_columns());
|
||||
assert_eq!("__primary_key", modified[0].schema().field(0).name());
|
||||
assert_eq!("greptime_timestamp", modified[0].schema().field(1).name());
|
||||
assert_eq!("greptime_value", modified[0].schema().field(2).name());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_logical_batches_to_physical_taxonomy_failure() {
|
||||
let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
|
||||
|
||||
let table_batches = vec![TableBatch {
|
||||
table_name: "t1".to_string(),
|
||||
table_id: 1,
|
||||
batches: vec![batch],
|
||||
row_count: 1,
|
||||
}];
|
||||
|
||||
// tag1 is missing from name_to_ids, causing columns_taxonomy to fail.
|
||||
let name_to_ids = HashMap::new();
|
||||
let partition_columns = HashSet::new();
|
||||
let err =
|
||||
transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("not found in physical table column IDs")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_logical_batches_to_physical_multiple_batches() {
|
||||
let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
|
||||
let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0);
|
||||
|
||||
let table_batches = vec![
|
||||
TableBatch {
|
||||
table_name: "t1".to_string(),
|
||||
table_id: 1,
|
||||
batches: vec![batch1],
|
||||
row_count: 1,
|
||||
},
|
||||
TableBatch {
|
||||
table_name: "t2".to_string(),
|
||||
table_id: 2,
|
||||
batches: vec![batch2],
|
||||
row_count: 1,
|
||||
},
|
||||
];
|
||||
|
||||
let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]);
|
||||
let partition_columns = HashSet::new();
|
||||
let modified =
|
||||
transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(2, modified.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transform_logical_batches_to_physical_mixed_success_failure() {
|
||||
let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
|
||||
let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0);
|
||||
|
||||
let table_batches = vec![
|
||||
TableBatch {
|
||||
table_name: "t1".to_string(),
|
||||
table_id: 1,
|
||||
batches: vec![batch1],
|
||||
row_count: 1,
|
||||
},
|
||||
TableBatch {
|
||||
table_name: "t2".to_string(),
|
||||
table_id: 2,
|
||||
batches: vec![batch2],
|
||||
row_count: 1,
|
||||
},
|
||||
];
|
||||
|
||||
// tag1 is missing from name_to_ids, causing batch1 to fail.
|
||||
let name_to_ids = HashMap::from([("tag2".to_string(), 2)]);
|
||||
let partition_columns = HashSet::new();
|
||||
let err =
|
||||
transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.to_string().contains("tag1"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,599 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use api::v1::ArrowIpc;
|
||||
use api::v1::meta::Peer;
|
||||
use api::v1::region::{
|
||||
BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
|
||||
};
|
||||
use arrow::compute::filter_record_batch;
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use async_trait::async_trait;
|
||||
use catalog::CatalogManagerRef;
|
||||
use common_grpc::error::Error as GrpcError;
|
||||
use common_grpc::flight::record_batch_to_ipc;
|
||||
use common_meta::node_manager::NodeManagerRef;
|
||||
use common_telemetry::tracing_context::TracingContext;
|
||||
use partition::manager::PartitionRuleManagerRef;
|
||||
use partition::partition::PartitionRuleRef;
|
||||
use snafu::ResultExt;
|
||||
use store_api::storage::RegionId;
|
||||
use table::metadata::{TableId, TableInfo, TableInfoRef};
|
||||
|
||||
use crate::batcher::logical_table::batch_convert::strip_partition_columns_from_batch;
|
||||
use crate::error;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PhysicalTableMetadata {
|
||||
pub table_info: TableInfoRef,
|
||||
/// Mapping from column name to column id
|
||||
pub col_name_to_ids: Option<HashMap<String, u32>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait PhysicalFlushCatalogProvider: Send + Sync {
|
||||
async fn physical_table(
|
||||
&self,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
table_name: &str,
|
||||
query_ctx: &session::context::QueryContext,
|
||||
) -> catalog::error::Result<Option<PhysicalTableMetadata>>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait PhysicalFlushPartitionProvider: Send + Sync {
|
||||
async fn find_table_partition_rule(
|
||||
&self,
|
||||
table_info: &TableInfo,
|
||||
) -> partition::error::Result<PartitionRuleRef>;
|
||||
|
||||
async fn find_region_leader(&self, region_id: RegionId) -> Result<Peer>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait PhysicalFlushNodeRequester: Send + Sync {
|
||||
async fn handle(
|
||||
&self,
|
||||
peer: &Peer,
|
||||
request: RegionRequest,
|
||||
) -> Result<api::region::RegionResponse>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(in crate::batcher::logical_table) struct CatalogManagerPhysicalFlushAdapter {
|
||||
pub(in crate::batcher::logical_table) catalog_manager: CatalogManagerRef,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl PhysicalFlushCatalogProvider for CatalogManagerPhysicalFlushAdapter {
|
||||
async fn physical_table(
|
||||
&self,
|
||||
catalog: &str,
|
||||
schema: &str,
|
||||
table_name: &str,
|
||||
query_ctx: &session::context::QueryContext,
|
||||
) -> catalog::error::Result<Option<PhysicalTableMetadata>> {
|
||||
self.catalog_manager
|
||||
.table(catalog, schema, table_name, Some(query_ctx))
|
||||
.await
|
||||
.map(|table| {
|
||||
table.map(|table| {
|
||||
let table_info = table.table_info();
|
||||
let name_to_ids = table_info.name_to_ids();
|
||||
PhysicalTableMetadata {
|
||||
table_info,
|
||||
col_name_to_ids: name_to_ids,
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(in crate::batcher::logical_table) struct PartitionManagerPhysicalFlushAdapter {
|
||||
pub(in crate::batcher::logical_table) partition_manager: PartitionRuleManagerRef,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter {
|
||||
async fn find_table_partition_rule(
|
||||
&self,
|
||||
table_info: &TableInfo,
|
||||
) -> partition::error::Result<PartitionRuleRef> {
|
||||
self.partition_manager
|
||||
.find_table_partition_rule(table_info)
|
||||
.await
|
||||
.map(|(rule, _)| rule)
|
||||
}
|
||||
|
||||
async fn find_region_leader(&self, region_id: RegionId) -> Result<Peer> {
|
||||
let peer = self.partition_manager.find_region_leader(region_id).await?;
|
||||
Ok(peer)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(in crate::batcher::logical_table) struct NodeManagerPhysicalFlushAdapter {
|
||||
pub(in crate::batcher::logical_table) node_manager: NodeManagerRef,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl PhysicalFlushNodeRequester for NodeManagerPhysicalFlushAdapter {
|
||||
async fn handle(
|
||||
&self,
|
||||
peer: &Peer,
|
||||
request: RegionRequest,
|
||||
) -> error::Result<api::region::RegionResponse> {
|
||||
let datanode = self.node_manager.datanode(peer).await;
|
||||
datanode
|
||||
.handle(request)
|
||||
.await
|
||||
.context(error::CommonMetaSnafu)
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) struct FlushRegionWrite {
|
||||
pub(in crate::batcher::logical_table) datanode: Peer,
|
||||
pub(in crate::batcher::logical_table) request: RegionRequest,
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) struct PlannedRegionBatch {
|
||||
pub(in crate::batcher::logical_table) region_id: RegionId,
|
||||
pub(in crate::batcher::logical_table) batch: RecordBatch,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl PlannedRegionBatch {
|
||||
pub(in crate::batcher::logical_table) fn num_rows(&self) -> usize {
|
||||
self.batch.num_rows()
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) struct ResolvedRegionBatch {
|
||||
pub(in crate::batcher::logical_table) planned: PlannedRegionBatch,
|
||||
pub(in crate::batcher::logical_table) datanode: Peer,
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn should_dispatch_concurrently(
|
||||
region_write_count: usize,
|
||||
) -> bool {
|
||||
region_write_count > 1
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) async fn flush_region_writes_concurrently(
|
||||
node_manager: &(impl PhysicalFlushNodeRequester + ?Sized),
|
||||
writes: Vec<FlushRegionWrite>,
|
||||
) -> Result<usize> {
|
||||
let mut affected_rows = 0;
|
||||
if !should_dispatch_concurrently(writes.len()) {
|
||||
for write in writes {
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_write_region"])
|
||||
.start_timer();
|
||||
affected_rows += node_manager
|
||||
.handle(&write.datanode, write.request)
|
||||
.await?
|
||||
.affected_rows;
|
||||
}
|
||||
return Ok(affected_rows);
|
||||
}
|
||||
|
||||
let write_futures = writes.into_iter().map(|write| async move {
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_write_region"])
|
||||
.start_timer();
|
||||
|
||||
let response = node_manager.handle(&write.datanode, write.request).await?;
|
||||
Ok::<_, Error>(response.affected_rows)
|
||||
});
|
||||
|
||||
// todo(hl): should be bounded.
|
||||
let affected_rows = futures::future::try_join_all(write_futures)
|
||||
.await?
|
||||
.into_iter()
|
||||
.sum();
|
||||
Ok(affected_rows)
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn split_combined_batch_by_region(
|
||||
combined_batch: &RecordBatch,
|
||||
partition_rule: &dyn partition::partition::PartitionRule,
|
||||
) -> Result<HashMap<u32, partition::partition::RegionMask>> {
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_split_record_batch"])
|
||||
.start_timer();
|
||||
let map = partition_rule.split_record_batch(combined_batch)?;
|
||||
Ok(map)
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn prepare_physical_region_routing_batch(
|
||||
combined_batch: RecordBatch,
|
||||
partition_columns: &[String],
|
||||
) -> Result<RecordBatch> {
|
||||
if partition_columns.is_empty() {
|
||||
return Ok(combined_batch);
|
||||
}
|
||||
strip_partition_columns_from_batch(combined_batch)
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn plan_region_batch(
|
||||
stripped_batch: &RecordBatch,
|
||||
physical_table_id: TableId,
|
||||
region_number: u32,
|
||||
mask: &partition::partition::RegionMask,
|
||||
) -> Result<Option<PlannedRegionBatch>> {
|
||||
if mask.select_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let region_batch = if mask.select_all() {
|
||||
stripped_batch.clone()
|
||||
} else {
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_filter_record_batch"])
|
||||
.start_timer();
|
||||
filter_record_batch(stripped_batch, mask.array()).context(error::ArrowSnafu)?
|
||||
};
|
||||
|
||||
let row_count = region_batch.num_rows();
|
||||
if row_count == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(PlannedRegionBatch {
|
||||
region_id: RegionId::new(physical_table_id, region_number),
|
||||
batch: region_batch,
|
||||
}))
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn plan_region_batches(
|
||||
combined_batch: RecordBatch,
|
||||
physical_table_id: TableId,
|
||||
partition_rule: &dyn partition::partition::PartitionRule,
|
||||
partition_columns: &[String],
|
||||
) -> Result<Vec<PlannedRegionBatch>> {
|
||||
let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?;
|
||||
let stripped_batch = prepare_physical_region_routing_batch(combined_batch, partition_columns)?;
|
||||
|
||||
let mut planned_batches = Vec::new();
|
||||
for (region_number, mask) in region_masks {
|
||||
if let Some(planned_batch) =
|
||||
plan_region_batch(&stripped_batch, physical_table_id, region_number, &mask)?
|
||||
{
|
||||
planned_batches.push(planned_batch);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(planned_batches)
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) async fn resolve_region_targets(
|
||||
planned_batches: Vec<PlannedRegionBatch>,
|
||||
partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
|
||||
) -> Result<Vec<ResolvedRegionBatch>> {
|
||||
let mut resolved_batches = Vec::with_capacity(planned_batches.len());
|
||||
for planned in planned_batches {
|
||||
let datanode = {
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_resolve_region_leader"])
|
||||
.start_timer();
|
||||
partition_manager
|
||||
.find_region_leader(planned.region_id)
|
||||
.await?
|
||||
};
|
||||
|
||||
resolved_batches.push(ResolvedRegionBatch { planned, datanode });
|
||||
}
|
||||
|
||||
Ok(resolved_batches)
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn encode_region_write_requests(
|
||||
resolved_batches: Vec<ResolvedRegionBatch>,
|
||||
skip_wal: bool,
|
||||
) -> Result<Vec<FlushRegionWrite>> {
|
||||
let mut region_writes = Vec::with_capacity(resolved_batches.len());
|
||||
for resolved in resolved_batches {
|
||||
let region_id = resolved.planned.region_id;
|
||||
let (schema_bytes, data_header, payload) = {
|
||||
let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
|
||||
.with_label_values(&["flush_physical_encode_ipc"])
|
||||
.start_timer();
|
||||
record_batch_to_ipc(resolved.planned.batch).map_err(map_ipc_error)?
|
||||
};
|
||||
|
||||
let request = RegionRequest {
|
||||
header: Some(RegionRequestHeader {
|
||||
tracing_context: TracingContext::from_current_span().to_w3c(),
|
||||
..Default::default()
|
||||
}),
|
||||
body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
|
||||
skip_wal,
|
||||
region_id: region_id.as_u64(),
|
||||
partition_expr_version: None,
|
||||
// Set aligned_schema_version to None so that datanode will check the batch schema again to see if any
|
||||
// column is missing.
|
||||
aligned_schema_version: None,
|
||||
body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
|
||||
schema: schema_bytes,
|
||||
data_header,
|
||||
payload,
|
||||
})),
|
||||
})),
|
||||
};
|
||||
|
||||
region_writes.push(FlushRegionWrite {
|
||||
datanode: resolved.datanode,
|
||||
request,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(region_writes)
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Error {
|
||||
match error {
|
||||
GrpcError::NotSupported { feat } => Error::NotSupported { feat },
|
||||
GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason },
|
||||
error => Error::Internal {
|
||||
err_msg: error.to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use api::region::RegionResponse;
|
||||
use api::v1::meta::Peer;
|
||||
use api::v1::region::{RegionRequest, region_request};
|
||||
use arrow::array::{BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray};
|
||||
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use async_trait::async_trait;
|
||||
use common_meta::error::Result as MetaResult;
|
||||
use common_meta::node_manager::{Datanode, DatanodeRef};
|
||||
use common_query::request::QueryRequest;
|
||||
use common_recordbatch::SendableRecordBatchStream;
|
||||
use partition::partition::{PartitionRule, RegionMask};
|
||||
use store_api::storage::RegionId;
|
||||
use tokio::time::sleep;
|
||||
|
||||
use crate::batcher::logical_table::region_write::{
|
||||
FlushRegionWrite, PlannedRegionBatch, ResolvedRegionBatch, encode_region_write_requests,
|
||||
flush_region_writes_concurrently, plan_region_batches, should_dispatch_concurrently,
|
||||
};
|
||||
use crate::batcher::logical_table::test_util::ConcurrentMockNodeManager;
|
||||
|
||||
struct TwoRegionPartitionRule {
|
||||
partition_columns: Vec<String>,
|
||||
}
|
||||
|
||||
impl PartitionRule for TwoRegionPartitionRule {
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
|
||||
fn partition_columns(&self) -> &[String] {
|
||||
&self.partition_columns
|
||||
}
|
||||
|
||||
fn find_region(
|
||||
&self,
|
||||
_values: &[datatypes::prelude::Value],
|
||||
) -> partition::error::Result<store_api::storage::RegionNumber> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn split_record_batch(
|
||||
&self,
|
||||
_record_batch: &RecordBatch,
|
||||
) -> partition::error::Result<HashMap<store_api::storage::RegionNumber, RegionMask>>
|
||||
{
|
||||
Ok(HashMap::from([
|
||||
(1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)),
|
||||
(2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)),
|
||||
(
|
||||
3,
|
||||
RegionMask::new(BooleanArray::from(vec![false, false]), 0),
|
||||
),
|
||||
]))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ConcurrentMockDatanode {
|
||||
delay: Duration,
|
||||
inflight: Arc<AtomicUsize>,
|
||||
max_inflight: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Datanode for ConcurrentMockDatanode {
|
||||
async fn handle(&self, _request: RegionRequest) -> MetaResult<RegionResponse> {
|
||||
let now = self.inflight.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
loop {
|
||||
let max = self.max_inflight.load(Ordering::SeqCst);
|
||||
if now <= max {
|
||||
break;
|
||||
}
|
||||
if self
|
||||
.max_inflight
|
||||
.compare_exchange(max, now, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_ok()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
sleep(self.delay).await;
|
||||
self.inflight.fetch_sub(1, Ordering::SeqCst);
|
||||
Ok(RegionResponse::new(0))
|
||||
}
|
||||
|
||||
async fn handle_query(
|
||||
&self,
|
||||
_request: QueryRequest,
|
||||
) -> MetaResult<SendableRecordBatchStream> {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_flush_region_writes_concurrently_dispatches_multiple_datanodes() {
|
||||
let inflight = Arc::new(AtomicUsize::new(0));
|
||||
let max_inflight = Arc::new(AtomicUsize::new(0));
|
||||
let datanode1: DatanodeRef = Arc::new(ConcurrentMockDatanode {
|
||||
delay: Duration::from_millis(100),
|
||||
inflight: inflight.clone(),
|
||||
max_inflight: max_inflight.clone(),
|
||||
});
|
||||
let datanode2: DatanodeRef = Arc::new(ConcurrentMockDatanode {
|
||||
delay: Duration::from_millis(100),
|
||||
inflight,
|
||||
max_inflight: max_inflight.clone(),
|
||||
});
|
||||
|
||||
let mut datanodes = HashMap::new();
|
||||
datanodes.insert(1, datanode1);
|
||||
datanodes.insert(2, datanode2);
|
||||
let node_manager = Arc::new(ConcurrentMockNodeManager {
|
||||
datanodes: Arc::new(datanodes),
|
||||
});
|
||||
|
||||
let writes = vec![
|
||||
FlushRegionWrite {
|
||||
datanode: Peer {
|
||||
id: 1,
|
||||
addr: "node1".to_string(),
|
||||
},
|
||||
request: RegionRequest::default(),
|
||||
},
|
||||
FlushRegionWrite {
|
||||
datanode: Peer {
|
||||
id: 2,
|
||||
addr: "node2".to_string(),
|
||||
},
|
||||
request: RegionRequest::default(),
|
||||
},
|
||||
];
|
||||
|
||||
flush_region_writes_concurrently(node_manager.as_ref(), writes)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(max_inflight.load(Ordering::SeqCst) >= 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_should_dispatch_concurrently_by_region_count() {
|
||||
assert!(!should_dispatch_concurrently(0));
|
||||
assert!(!should_dispatch_concurrently(1));
|
||||
assert!(should_dispatch_concurrently(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_plan_region_batches_splits_and_strips_partition_columns() {
|
||||
let combined_batch = RecordBatch::try_new(
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("__primary_key", ArrowDataType::Binary, false),
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new("host", ArrowDataType::Utf8, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(BinaryArray::from(vec![b"k1".as_slice(), b"k2".as_slice()])),
|
||||
Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000_i64])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![1.0_f64, 2.0_f64])),
|
||||
Arc::new(StringArray::from(vec!["node-1", "node-2"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let mut planned_batches = plan_region_batches(
|
||||
combined_batch,
|
||||
1024,
|
||||
&TwoRegionPartitionRule {
|
||||
partition_columns: vec!["host".to_string()],
|
||||
},
|
||||
&["host".to_string()],
|
||||
)
|
||||
.unwrap();
|
||||
planned_batches.sort_by_key(|planned| planned.region_id.region_number());
|
||||
|
||||
assert_eq!(2, planned_batches.len());
|
||||
assert_eq!(RegionId::new(1024, 1), planned_batches[0].region_id);
|
||||
assert_eq!(1, planned_batches[0].num_rows());
|
||||
assert_eq!(3, planned_batches[0].batch.num_columns());
|
||||
assert_eq!(RegionId::new(1024, 2), planned_batches[1].region_id);
|
||||
assert_eq!(1, planned_batches[1].num_rows());
|
||||
assert_eq!(3, planned_batches[1].batch.num_columns());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encode_region_write_requests_builds_bulk_insert_requests() {
|
||||
check_encode_region_write_requests(false);
|
||||
check_encode_region_write_requests(true);
|
||||
}
|
||||
|
||||
fn check_encode_region_write_requests(skip_wal: bool) {
|
||||
let planned_batch = PlannedRegionBatch {
|
||||
region_id: RegionId::new(1024, 1),
|
||||
batch: RecordBatch::try_new(
|
||||
Arc::new(ArrowSchema::new(vec![
|
||||
Field::new("__primary_key", ArrowDataType::Binary, false),
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
|
||||
Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![1.0_f64])),
|
||||
],
|
||||
)
|
||||
.unwrap(),
|
||||
};
|
||||
let resolved_batch = ResolvedRegionBatch {
|
||||
planned: planned_batch,
|
||||
datanode: Peer {
|
||||
id: 1,
|
||||
addr: "node-1".to_string(),
|
||||
},
|
||||
};
|
||||
let writes = encode_region_write_requests(vec![resolved_batch], skip_wal).unwrap();
|
||||
|
||||
assert_eq!(1, writes.len());
|
||||
assert_eq!(1, writes[0].datanode.id);
|
||||
let Some(region_request::Body::BulkInsert(request)) = &writes[0].request.body else {
|
||||
panic!("expected bulk insert request");
|
||||
};
|
||||
assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id);
|
||||
assert_eq!(skip_wal, request.skip_wal);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
// Copyright 2023 Greptime Team
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use api::region::RegionResponse;
|
||||
use api::v1::flow::{DirtyWindowRequests, FlowRequest, FlowResponse};
|
||||
use api::v1::meta::Peer;
|
||||
use api::v1::region::{InsertRequests, RegionRequest};
|
||||
use arrow::array::{StringArray, TimestampMillisecondArray};
|
||||
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use async_trait::async_trait;
|
||||
use common_meta::cache::{TableFlownodeSetCacheRef, new_table_flownode_set_cache};
|
||||
use common_meta::error::Result as MetaResult;
|
||||
use common_meta::instruction::{CacheIdent, CreateFlow};
|
||||
use common_meta::kv_backend::memory::MemoryKvBackend;
|
||||
use common_meta::node_manager::{
|
||||
DatanodeManager, DatanodeRef, Flownode, FlownodeManager, FlownodeRef,
|
||||
};
|
||||
use moka::future::CacheBuilder;
|
||||
use snafu::ResultExt;
|
||||
use table::metadata::TableId;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::batcher::logical_table::batch_convert::RecordBatchWithTsIdx;
|
||||
use crate::batcher::logical_table::region_write::PhysicalFlushNodeRequester;
|
||||
use crate::error;
|
||||
|
||||
pub(in crate::batcher::logical_table) fn mock_tag_batch(
|
||||
tag_name: &str,
|
||||
tag_value: &str,
|
||||
ts: i64,
|
||||
val: f64,
|
||||
) -> RecordBatch {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
Field::new(
|
||||
"greptime_timestamp",
|
||||
ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
|
||||
false,
|
||||
),
|
||||
Field::new("greptime_value", ArrowDataType::Float64, true),
|
||||
Field::new(tag_name, ArrowDataType::Utf8, true),
|
||||
]));
|
||||
|
||||
RecordBatch::try_new(
|
||||
schema,
|
||||
vec![
|
||||
Arc::new(TimestampMillisecondArray::from(vec![ts])),
|
||||
Arc::new(arrow::array::Float64Array::from(vec![val])),
|
||||
Arc::new(StringArray::from(vec![tag_value])),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) fn mock_aligned_tag_batch(
|
||||
tag_name: &str,
|
||||
tag_value: &str,
|
||||
ts: i64,
|
||||
val: f64,
|
||||
) -> RecordBatchWithTsIdx {
|
||||
RecordBatchWithTsIdx::try_new(mock_tag_batch(tag_name, tag_value, ts, val), 0).unwrap()
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(in crate::batcher::logical_table) struct ConcurrentMockNodeManager {
|
||||
pub(in crate::batcher::logical_table) datanodes: Arc<HashMap<u64, DatanodeRef>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DatanodeManager for ConcurrentMockNodeManager {
|
||||
async fn datanode(&self, node: &Peer) -> DatanodeRef {
|
||||
self.datanodes
|
||||
.get(&node.id)
|
||||
.expect("datanode not found")
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) struct NoopFlownode;
|
||||
|
||||
#[async_trait]
|
||||
impl Flownode for NoopFlownode {
|
||||
async fn handle(&self, _request: FlowRequest) -> MetaResult<FlowResponse> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
async fn handle_inserts(&self, _request: InsertRequests) -> MetaResult<FlowResponse> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
async fn handle_mark_window_dirty(
|
||||
&self,
|
||||
_req: DirtyWindowRequests,
|
||||
) -> MetaResult<FlowResponse> {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl FlownodeManager for ConcurrentMockNodeManager {
|
||||
async fn flownode(&self, _node: &Peer) -> FlownodeRef {
|
||||
Arc::new(NoopFlownode)
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) struct RecordingFlownode {
|
||||
pub(in crate::batcher::logical_table) requests_tx: mpsc::UnboundedSender<DirtyWindowRequests>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Flownode for RecordingFlownode {
|
||||
async fn handle(&self, _request: FlowRequest) -> MetaResult<FlowResponse> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
async fn handle_inserts(&self, _request: InsertRequests) -> MetaResult<FlowResponse> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
async fn handle_mark_window_dirty(&self, req: DirtyWindowRequests) -> MetaResult<FlowResponse> {
|
||||
self.requests_tx.send(req).unwrap();
|
||||
Ok(FlowResponse::default())
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) struct FlowNotificationMockNodeManager {
|
||||
pub(in crate::batcher::logical_table) flownode: FlownodeRef,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DatanodeManager for FlowNotificationMockNodeManager {
|
||||
async fn datanode(&self, _node: &Peer) -> DatanodeRef {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl FlownodeManager for FlowNotificationMockNodeManager {
|
||||
async fn flownode(&self, _node: &Peer) -> FlownodeRef {
|
||||
self.flownode.clone()
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::batcher::logical_table) async fn mock_table_flownode_cache(
|
||||
table_id: TableId,
|
||||
peer: Peer,
|
||||
) -> TableFlownodeSetCacheRef {
|
||||
let cache = Arc::new(new_table_flownode_set_cache(
|
||||
"test".to_string(),
|
||||
CacheBuilder::new(1).build(),
|
||||
Arc::new(MemoryKvBackend::default()),
|
||||
));
|
||||
cache
|
||||
.invalidate(&[CacheIdent::CreateFlow(CreateFlow {
|
||||
flow_id: 1,
|
||||
source_table_ids: vec![table_id],
|
||||
partition_to_peer_mapping: vec![(0, peer.clone()), (1, peer)],
|
||||
})])
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl PhysicalFlushNodeRequester for ConcurrentMockNodeManager {
|
||||
async fn handle(&self, peer: &Peer, request: RegionRequest) -> error::Result<RegionResponse> {
|
||||
let datanode = self.datanode(peer).await;
|
||||
datanode
|
||||
.handle(request)
|
||||
.await
|
||||
.context(error::CommonMetaSnafu)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user