feat(metric-engine): add integer value companions

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2026-07-10 15:40:02 +08:00
parent d28bca855f
commit a18ef07f1e
17 changed files with 1238 additions and 240 deletions
Generated
+1
View File
@@ -8332,6 +8332,7 @@ dependencies = [
"common-wal",
"criterion 0.7.0",
"datafusion",
"datafusion-physical-expr",
"datatypes",
"futures-util",
"fxhash",
+1
View File
@@ -26,6 +26,7 @@ common-telemetry.workspace = true
common-time.workspace = true
chrono.workspace = true
datafusion.workspace = true
datafusion-physical-expr.workspace = true
datatypes.workspace = true
futures-util.workspace = true
humantime-serde.workspace = true
+4 -1
View File
@@ -312,11 +312,14 @@ mod test {
let column_names = new_metadata
.column_metadatas
.iter()
.map(|c| &c.column_schema.name)
.map(|c| c.column_schema.name.as_str())
.collect::<Vec<_>>();
let value_int_name =
store_api::metric_engine_consts::metric_engine_value_int_column_name(greptime_value());
let expected = vec![
greptime_timestamp(),
greptime_value(),
value_int_name.as_str(),
"__table_id",
"__tsid",
"job",
+11 -5
View File
@@ -279,7 +279,9 @@ mod test {
use common_meta::ddl::test_util::assert_column_name_and_id;
use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions};
use common_query::prelude::{greptime_timestamp, greptime_value};
use store_api::metric_engine_consts::ALTER_PHYSICAL_EXTENSION_KEY;
use store_api::metric_engine_consts::{
ALTER_PHYSICAL_EXTENSION_KEY, metric_engine_value_int_column_name,
};
use store_api::region_engine::RegionEngine;
use store_api::region_request::{
AlterKind, BatchRegionDdlRequest, RegionAlterRequest, SetRegionOption,
@@ -357,15 +359,17 @@ mod test {
assert_eq!(timestamp_index, SemanticType::Timestamp);
let column_metadatas =
parse_column_metadatas(&response.extensions, ALTER_PHYSICAL_EXTENSION_KEY).unwrap();
let value_int_name = metric_engine_value_int_column_name(greptime_value());
assert_column_name_and_id(
&column_metadatas,
&[
(greptime_timestamp(), 0),
(greptime_value(), 1),
(value_int_name.as_str(), 2),
("__table_id", ReservedColumnId::table_id()),
("__tsid", ReservedColumnId::tsid()),
("job", 2),
("tag1", 3),
("job", 3),
("tag1", 4),
],
);
}
@@ -416,15 +420,17 @@ mod test {
let column_metadatas =
parse_column_metadatas(&response.extensions, ALTER_PHYSICAL_EXTENSION_KEY).unwrap();
let value_int_name = metric_engine_value_int_column_name(greptime_value());
assert_column_name_and_id(
&column_metadatas,
&[
(greptime_timestamp(), 0),
(greptime_value(), 1),
(value_int_name.as_str(), 2),
("__table_id", ReservedColumnId::table_id()),
("__tsid", ReservedColumnId::tsid()),
("job", 2),
("tag1", 3),
("job", 3),
("tag1", 4),
],
);
}
+120 -1
View File
@@ -13,17 +13,22 @@
// limitations under the License.
use std::collections::HashSet;
use std::sync::Arc;
use api::v1::{ArrowIpc, SemanticType};
use bytes::Bytes;
use common_grpc::flight::{FlightEncoder, FlightMessage};
use datatypes::arrow::array::new_null_array;
use datatypes::arrow::datatypes::{Field as ArrowField, Schema as ArrowSchema};
use datatypes::arrow::record_batch::RecordBatch;
use snafu::{OptionExt, ResultExt, ensure};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::RegionMetadataRef;
use store_api::metric_engine_consts::is_metric_engine_value_int_column;
use store_api::region_engine::RegionEngine;
use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest};
use store_api::storage::RegionId;
use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
use crate::batch_modifier::{TagColumnInfo, modify_batch_sparse};
use crate::engine::MetricEngineInner;
@@ -74,6 +79,15 @@ impl MetricEngineInner {
region_id: RegionId,
mut request: RegionBulkInsertsRequest,
) -> Result<AffectedRows> {
request.payload = self
.append_missing_metric_value_int_fields(region_id, request.payload)
.await?;
let (schema, data_header, payload) = record_batch_to_ipc(&request.payload)?;
request.raw_data = ArrowIpc {
schema,
data_header,
payload,
};
// Simply set the aligned schema to the data region schema version to avoid filling missing columns
// because that schema should be constant and callers have ensured request has the same schema.
request.aligned_schema_version = Some(self.physical_schema_version(region_id).await?);
@@ -118,6 +132,9 @@ impl MetricEngineInner {
&tag_columns,
&non_tag_indices,
)?;
let modified_batch = self
.append_missing_metric_value_int_fields(data_region_id, modified_batch)
.await?;
let (schema, data_header, payload) = record_batch_to_ipc(&modified_batch)?;
let partition_expr_version = request.partition_expr_version;
@@ -139,6 +156,69 @@ impl MetricEngineInner {
.await
}
async fn append_missing_metric_value_int_fields(
&self,
data_region_id: RegionId,
batch: RecordBatch,
) -> Result<RecordBatch> {
let metadata = self
.mito
.get_metadata(data_region_id)
.await
.context(error::MitoReadOperationSnafu)?;
let batch_schema = batch.schema();
let existing_names = batch_schema
.fields()
.iter()
.map(|field| field.name().as_str())
.collect::<HashSet<_>>();
let missing_columns = metadata
.field_columns()
.filter(|column| is_metric_engine_value_int_column(&column.column_schema.name))
.filter(|column| !existing_names.contains(column.column_schema.name.as_str()))
.collect::<Vec<_>>();
if missing_columns.is_empty() {
return Ok(batch);
}
let mut fields = batch
.schema()
.fields()
.iter()
.map(|field| field.as_ref().clone())
.collect::<Vec<_>>();
let insert_index = batch_schema
.index_of(PRIMARY_KEY_COLUMN_NAME)
.unwrap_or(fields.len());
let mut columns = batch.columns().to_vec();
for (offset, column) in missing_columns.into_iter().enumerate() {
let field = ArrowField::try_from(&column.column_schema).map_err(|err| {
error::InvalidRequestSnafu {
region_id: data_region_id,
reason: format!(
"failed to build Arrow field for column {}: {err}",
column.column_schema.name
),
}
.build()
})?;
columns.insert(
insert_index + offset,
new_null_array(field.data_type(), batch.num_rows()),
);
fields.insert(insert_index + offset, field);
}
RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).map_err(|err| {
error::InvalidRequestSnafu {
region_id: data_region_id,
reason: format!("failed to append metric value companion columns: {err}"),
}
.build()
})
}
async fn physical_schema_version(&self, region_id: RegionId) -> Result<u64> {
Ok(self
.mito
@@ -302,6 +382,25 @@ mod tests {
})
}
fn collect_metric_values(batches: &RecordBatches) -> Vec<f64> {
let mut values = batches
.iter()
.flat_map(|batch| {
batch
.column_by_name(greptime_value())
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
.iter()
.map(Option::unwrap)
.collect::<Vec<_>>()
})
.collect::<Vec<_>>();
values.sort_by(f64::total_cmp);
values
}
async fn init_dense_metric_region(env: &TestEnv) -> RegionId {
let physical_region_id = env.default_physical_region_id();
env.create_physical_region(
@@ -374,7 +473,7 @@ mod tests {
let tag_columns = vec![TagColumnInfo {
name: "job".to_string(),
index: 2,
column_id: 2, // column_id for "job" in the physical table
column_id: 3, // column_id for "job" in the physical table
}];
let non_tag_indices = vec![0, 1]; // timestamp, value
let second_batch = build_logical_batch(3, 3);
@@ -401,6 +500,26 @@ mod tests {
.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 6);
assert_eq!(
collect_metric_values(&batches),
vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
);
env.metric()
.handle_request(physical_region_id, RegionRequest::Flush(Default::default()))
.await
.unwrap();
let stream = env
.metric()
.scan_to_stream(logical_region_id, ScanRequest::default())
.await
.unwrap();
let batches = RecordBatches::try_collect(stream).await.unwrap();
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 6);
assert_eq!(
collect_metric_values(&batches),
vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
);
}
#[tokio::test]
+162 -16
View File
@@ -28,11 +28,12 @@ use snafu::{OptionExt, ResultExt, ensure};
use store_api::metadata::ColumnMetadata;
use store_api::metric_engine_consts::{
ALTER_PHYSICAL_EXTENSION_KEY, DATA_REGION_SUBDIR, DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, METADATA_REGION_SUBDIR,
METADATA_SCHEMA_KEY_COLUMN_INDEX, METADATA_SCHEMA_KEY_COLUMN_NAME,
DATA_SCHEMA_TSID_COLUMN_NAME, DATA_SCHEMA_VALUE_INT_COLUMN_SUFFIX, LOGICAL_TABLE_METADATA_KEY,
METADATA_REGION_SUBDIR, METADATA_SCHEMA_KEY_COLUMN_INDEX, METADATA_SCHEMA_KEY_COLUMN_NAME,
METADATA_SCHEMA_TIMESTAMP_COLUMN_INDEX, METADATA_SCHEMA_TIMESTAMP_COLUMN_NAME,
METADATA_SCHEMA_VALUE_COLUMN_INDEX, METADATA_SCHEMA_VALUE_COLUMN_NAME,
is_metric_engine_internal_column,
is_metric_engine_internal_column, is_metric_engine_value_int_column,
metric_engine_value_int_column_name,
};
use store_api::mito_engine_options::{TTL_KEY, WAL_OPTIONS_KEY};
use store_api::region_engine::RegionEngine;
@@ -302,19 +303,35 @@ impl MetricEngineInner {
.iter()
.map(|(region_id, _)| *region_id)
.collect::<Vec<_>>();
let logical_region_columns = requests.iter().map(|(region_id, request)| {
let logical_column_metadata = requests
.iter()
.map(|(region_id, request)| {
(
*region_id,
request
.column_metadatas
.iter()
.map(|metadata| {
// Safety: previous steps ensure the physical region exist
let physical_metadata = *physical_schema_map
.get(metadata.column_schema.name.as_str())
.unwrap();
let mut column_metadata = physical_metadata.clone();
if metadata.semantic_type == SemanticType::Field {
column_metadata.column_schema = metadata.column_schema.clone();
}
(metadata.column_schema.name.clone(), column_metadata)
})
.collect::<HashMap<_, _>>(),
)
})
.collect::<Vec<_>>();
let logical_region_columns = logical_column_metadata.iter().map(|(region_id, columns)| {
(
*region_id,
request
.column_metadatas
columns
.iter()
.map(|metadata| {
// Safety: previous steps ensure the physical region exist
let column_metadata = *physical_schema_map
.get(metadata.column_schema.name.as_str())
.unwrap();
(metadata.column_schema.name.as_str(), column_metadata)
})
.map(|(name, column_metadata)| (name.as_str(), column_metadata))
.collect::<HashMap<_, _>>(),
)
});
@@ -379,6 +396,13 @@ impl MetricEngineInner {
column: DATA_SCHEMA_TSID_COLUMN_NAME,
}
);
for name in name_to_index.keys() {
ensure!(
!is_metric_engine_value_int_column(name)
|| is_valid_physical_metric_value_int_column(request, name),
InternalColumnOccupiedSnafu { column: name }
);
}
// check if required table option is present
ensure!(
@@ -529,6 +553,7 @@ impl MetricEngineInner {
let table_id_col_def = request.column_metadatas.iter().any(is_metric_name_col);
let tsid_col_def = request.column_metadatas.iter().any(is_tsid_col);
append_metric_value_int_columns(&mut data_region_request.column_metadatas);
// change nullability for tag columns
data_region_request
@@ -563,6 +588,80 @@ impl MetricEngineInner {
}
}
fn is_valid_physical_metric_value_int_column(request: &RegionCreateRequest, name: &str) -> bool {
if !request.is_physical_table() {
return false;
}
let Some(value_name) = name.strip_suffix(DATA_SCHEMA_VALUE_INT_COLUMN_SUFFIX) else {
return false;
};
let find_column = |name| {
request
.column_metadatas
.iter()
.find(|metadata| metadata.column_schema.name == name)
};
let Some(value_column) = find_column(value_name) else {
return false;
};
let Some(int_column) = find_column(name) else {
return false;
};
value_column.semantic_type == SemanticType::Field
&& value_column.column_schema.data_type == ConcreteDataType::float64_datatype()
&& int_column.semantic_type == SemanticType::Field
&& int_column.column_schema.data_type == ConcreteDataType::int64_datatype()
}
fn append_metric_value_int_columns(column_metadatas: &mut Vec<ColumnMetadata>) {
let mut next_column_id = column_metadatas
.iter()
.map(|metadata| metadata.column_id)
.filter(|column_id| !ReservedColumnId::is_reserved(*column_id))
.max()
.unwrap_or(0)
+ 1;
let existing_names = column_metadatas
.iter()
.map(|metadata| metadata.column_schema.name.clone())
.collect::<HashSet<_>>();
let mut int_columns = Vec::new();
for metadata in column_metadatas.iter_mut() {
if is_metric_engine_value_int_column(&metadata.column_schema.name) {
metadata.column_schema.set_nullable();
continue;
}
if metadata.semantic_type != SemanticType::Field
|| metadata.column_schema.data_type != ConcreteDataType::float64_datatype()
{
continue;
}
metadata.column_schema.set_nullable();
let int_column_name = metric_engine_value_int_column_name(&metadata.column_schema.name);
if existing_names.contains(&int_column_name) {
continue;
}
int_columns.push(ColumnMetadata {
column_id: next_column_id,
semantic_type: SemanticType::Field,
column_schema: ColumnSchema::new(
int_column_name,
ConcreteDataType::int64_datatype(),
true,
),
});
next_column_id += 1;
}
column_metadatas.extend(int_columns);
}
fn table_id_col() -> ColumnMetadata {
ColumnMetadata {
column_id: ReservedColumnId::table_id(),
@@ -721,7 +820,8 @@ mod test {
);
// allow reserved internal columns when defined properly
let request = RegionCreateRequest {
let value_int_name = metric_engine_value_int_column_name("column2");
let mut request = RegionCreateRequest {
column_metadatas: vec![
ColumnMetadata {
column_id: 0,
@@ -750,6 +850,15 @@ mod test {
false,
),
},
ColumnMetadata {
column_id: 3,
semantic_type: SemanticType::Field,
column_schema: ColumnSchema::new(
&value_int_name,
ConcreteDataType::int64_datatype(),
false,
),
},
table_id_col(),
tsid_col(),
],
@@ -764,6 +873,21 @@ mod test {
requirements: Default::default(),
};
MetricEngineInner::verify_region_create_request(&request).unwrap();
append_metric_value_int_columns(&mut request.column_metadatas);
assert!(request.column_metadatas[2].column_schema.is_nullable());
assert!(request.column_metadatas[3].column_schema.is_nullable());
request.options = [(LOGICAL_TABLE_METADATA_KEY.to_string(), String::new())]
.into_iter()
.collect();
assert!(MetricEngineInner::verify_region_create_request(&request).is_err());
request.options = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
.into_iter()
.collect();
request.column_metadatas[3].column_schema =
ColumnSchema::new(&value_int_name, ConcreteDataType::string_datatype(), true);
assert!(MetricEngineInner::verify_region_create_request(&request).is_err());
// valid request
let request = RegionCreateRequest {
@@ -1074,7 +1198,7 @@ mod test {
let engine_inner = engine.inner;
let data_region_request = engine_inner.create_request_for_data_region(&request);
assert_eq!(data_region_request.column_metadatas.len(), 5);
assert_eq!(data_region_request.column_metadatas.len(), 6);
assert_eq!(
data_region_request.primary_key,
vec![ReservedColumnId::table_id(), ReservedColumnId::tsid(), 1]
@@ -1100,6 +1224,26 @@ mod test {
.unwrap();
assert!(tag_metadata.column_schema.is_nullable());
let value_metadata = data_region_request
.column_metadatas
.iter()
.find(|metadata| metadata.column_schema.name == "value")
.unwrap();
assert!(value_metadata.column_schema.is_nullable());
let value_int_name = metric_engine_value_int_column_name("value");
let value_int_metadata = data_region_request
.column_metadatas
.iter()
.find(|metadata| metadata.column_schema.name == value_int_name)
.unwrap();
assert_eq!(value_int_metadata.column_id, 3);
assert_eq!(
value_int_metadata.column_schema.data_type,
ConcreteDataType::int64_datatype()
);
assert!(value_int_metadata.column_schema.is_nullable());
let table_id_metadata = data_region_request
.column_metadatas
.iter()
@@ -1149,14 +1293,16 @@ mod test {
let column_metadatas =
parse_column_metadatas(&response.extensions, ALTER_PHYSICAL_EXTENSION_KEY).unwrap();
let value_int_name = metric_engine_value_int_column_name(greptime_value());
assert_column_name_and_id(
&column_metadatas,
&[
(greptime_timestamp(), 0),
(greptime_value(), 1),
(value_int_name.as_str(), 2),
("__table_id", ReservedColumnId::table_id()),
("__tsid", ReservedColumnId::tsid()),
("job", 2),
("job", 3),
],
);
}
+3 -3
View File
@@ -133,11 +133,11 @@ mod tests {
assert_eq!(
debug_format,
r#"
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/<file_id>.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000001/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/<file_id>.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000002/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/<file_id>.parquet", file_size: 4000, index_file_path: Some("test_metric_region/11_0000000001/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x03\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x03\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/<file_id>.parquet", file_size: 4000, index_file_path: Some("test_metric_region/11_0000000002/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x03\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x03\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/<file_id>.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/<file_id>.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/<file_id>.parquet", file_size: 3000, index_file_path: Some("test_metric_region/22_0000000042/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/<file_id>.parquet", file_size: 4000, index_file_path: Some("test_metric_region/22_0000000042/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x03\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x03\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/<file_id>.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#,
);
// list from storage
+171
View File
@@ -770,6 +770,10 @@ mod tests {
use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD;
use common_query::prelude::{greptime_timestamp, greptime_value};
use common_recordbatch::RecordBatches;
use datafusion::logical_expr::{col as df_col, lit as df_lit};
use datatypes::arrow::array::{
Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray,
};
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema};
use datatypes::value::Value as PartitionValue;
@@ -875,6 +879,50 @@ mod tests {
.unwrap()
}
type ValueRow = (String, i64, Option<f64>, Option<i64>);
fn collect_value_rows(batches: &RecordBatches, int_column_name: Option<&str>) -> Vec<ValueRow> {
let mut rows = Vec::new();
for batch in batches.iter() {
let jobs = batch
.column_by_name("job")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let values = batch
.column_by_name(greptime_value())
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
let timestamps = batch
.column_by_name(greptime_timestamp())
.unwrap()
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap();
let int_values = int_column_name.map(|name| {
batch
.column_by_name(name)
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
});
for row in 0..batch.num_rows() {
rows.push((
jobs.value(row).to_string(),
timestamps.value(row),
(!values.is_null(row)).then(|| values.value(row)),
int_values.and_then(|values| (!values.is_null(row)).then(|| values.value(row))),
));
}
}
rows.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
rows
}
async fn run_batch_write_with_schema_variants(
env: &TestEnv,
physical_region_id: RegionId,
@@ -1152,6 +1200,129 @@ mod tests {
assert_eq!(expected, batches.pretty_print().unwrap(), "logical region");
}
#[tokio::test]
async fn test_metric_value_split_roundtrip_after_flush() {
let env = TestEnv::new().await;
env.init_metric_region().await;
let schema = test_util::row_schema_with_tags(&["job"]);
let rows = [
(0, 1.0, "integer"),
(1, 2.0, "integer"),
(0, 1.5, "float"),
(1, 2.0, "float"),
]
.into_iter()
.map(|(timestamp, value, job)| Row {
values: vec![
ValueData::TimestampMillisecondValue(timestamp).into(),
ValueData::F64Value(value).into(),
ValueData::StringValue(job.to_string()).into(),
],
})
.collect();
let logical_region_id = env.default_logical_region_id();
env.metric()
.handle_request(
logical_region_id,
RegionRequest::Put(RegionPutRequest {
rows: Rows { schema, rows },
hint: None,
partition_expr_version: None,
}),
)
.await
.unwrap();
env.metric()
.handle_request(
env.default_physical_region_id(),
RegionRequest::Flush(Default::default()),
)
.await
.unwrap();
let physical_region_id = env.default_physical_region_id();
let visible_physical_stream = env
.metric()
.scan_to_stream(physical_region_id, ScanRequest::default())
.await
.unwrap();
let int_column_name =
store_api::metric_engine_consts::metric_engine_value_int_column_name(greptime_value());
assert!(
visible_physical_stream
.schema()
.column_schema_by_name(&int_column_name)
.is_none(),
"metric physical reads should hide split companion column"
);
let raw_physical_batches = RecordBatches::try_collect(
env.mito()
.scan_to_stream(
to_data_region_id(physical_region_id),
ScanRequest::default(),
)
.await
.unwrap(),
)
.await
.unwrap();
let physical_rows = collect_value_rows(&raw_physical_batches, Some(&int_column_name));
assert_eq!(
physical_rows,
vec![
("float".to_string(), 0, Some(1.5), None),
("float".to_string(), 1, Some(2.0), None),
("integer".to_string(), 0, None, Some(1)),
("integer".to_string(), 1, None, Some(2)),
]
);
let logical_batches = RecordBatches::try_collect(
env.metric()
.scan_to_stream(logical_region_id, ScanRequest::default())
.await
.unwrap(),
)
.await
.unwrap();
let logical_rows = collect_value_rows(&logical_batches, None);
assert_eq!(
logical_rows,
vec![
("float".to_string(), 0, Some(1.5), None),
("float".to_string(), 1, Some(2.0), None),
("integer".to_string(), 0, Some(1.0), None),
("integer".to_string(), 1, Some(2.0), None),
]
);
let filtered_batches = RecordBatches::try_collect(
env.metric()
.scan_to_stream(
logical_region_id,
ScanRequest {
filters: vec![df_col(greptime_value()).gt(df_lit(1.5_f64))],
..Default::default()
},
)
.await
.unwrap(),
)
.await
.unwrap();
let filtered_rows = collect_value_rows(&filtered_batches, None);
assert_eq!(
filtered_rows,
vec![
("float".to_string(), 1, Some(2.0), None),
("integer".to_string(), 1, Some(2.0), None),
]
);
}
#[tokio::test]
async fn test_write_logical_region_row_count() {
let env = TestEnv::new().await;
+334 -115
View File
@@ -12,23 +12,37 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::sync::Arc;
use api::v1::SemanticType;
use common_telemetry::{debug, error, tracing};
use datafusion::arrow::datatypes::DataType as ArrowDataType;
use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
use datafusion::common::{Column, Result as DataFusionResult};
use datafusion::functions::expr_fn::coalesce;
use datafusion::logical_expr::expr_fn::cast;
use datafusion::logical_expr::utils::expr_to_columns;
use datafusion::logical_expr::{self, Expr};
use datatypes::prelude::ConcreteDataType;
use snafu::{OptionExt, ResultExt};
use store_api::metadata::{RegionMetadataBuilder, RegionMetadataRef};
use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME;
use store_api::metric_engine_consts::{
DATA_SCHEMA_TABLE_ID_COLUMN_NAME, is_metric_engine_value_int_column,
metric_engine_value_int_column_name,
};
use store_api::region_engine::{RegionEngine, RegionScannerRef};
use store_api::region_request::AlterKind;
use store_api::storage::{RegionId, ScanRequest, SequenceNumber};
use crate::engine::MetricEngineInner;
use crate::error::{
InvalidMetadataSnafu, LogicalRegionNotFoundSnafu, MitoReadOperationSnafu, Result,
InvalidMetadataSnafu, InvalidRequestSnafu, LogicalRegionNotFoundSnafu, MitoReadOperationSnafu,
Result,
};
use crate::metrics::MITO_OPERATION_ELAPSED;
use crate::utils;
use crate::value_split::{ValueColumnProjection, ValueSplitProjectionMapper, ValueSplitScanner};
impl MetricEngineInner {
#[tracing::instrument(skip_all)]
@@ -59,10 +73,32 @@ impl MetricEngineInner {
.with_label_values(&["read_physical"])
.start_timer();
self.mito
.handle_query(region_id, request)
let data_region_id = utils::to_data_region_id(region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
.await
.context(MitoReadOperationSnafu)
.context(MitoReadOperationSnafu)?;
let visible_metadata = visible_physical_region_metadata(&physical_metadata)?;
let (request, mapper) = self.transform_request_with_mapper(
data_region_id,
request,
&visible_metadata,
&physical_metadata,
None,
)?;
let scanner = self
.mito
.handle_query(data_region_id, request)
.await
.context(MitoReadOperationSnafu)?;
Ok(Box::new(ValueSplitScanner::new(
scanner,
visible_metadata,
mapper,
)))
}
async fn read_logical_region(
@@ -76,8 +112,16 @@ impl MetricEngineInner {
let physical_region_id = self.get_physical_region_id(logical_region_id).await?;
let data_region_id = utils::to_data_region_id(physical_region_id);
let request = self
.transform_request(physical_region_id, logical_region_id, request)
let logical_metadata = self
.logical_region_metadata(physical_region_id, logical_region_id)
.await?;
let (request, mapper) = self
.transform_logical_request_with_mapper(
physical_region_id,
logical_region_id,
request,
&logical_metadata,
)
.await?;
let mut scanner = self
.mito
@@ -87,7 +131,11 @@ impl MetricEngineInner {
scanner.set_logical_region(true);
scanner.set_query_load_region_id(data_region_id);
Ok(scanner)
Ok(Box::new(ValueSplitScanner::new(
scanner,
logical_metadata,
mapper,
)))
}
pub async fn get_last_seq_num(&self, region_id: RegionId) -> Result<SequenceNumber> {
@@ -108,10 +156,13 @@ impl MetricEngineInner {
self.state.read().unwrap().exist_physical_region(region_id);
if is_reading_physical_region {
self.mito
.get_metadata(region_id)
let data_region_id = utils::to_data_region_id(region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
.await
.context(MitoReadOperationSnafu)
.context(MitoReadOperationSnafu)?;
visible_physical_region_metadata(&physical_metadata)
} else {
let physical_region_id = self.get_physical_region_id(region_id).await?;
self.logical_region_metadata(physical_region_id, region_id)
@@ -137,103 +188,162 @@ impl MetricEngineInner {
}
/// Transform the [ScanRequest] from logical region to physical data region.
#[cfg(test)]
async fn transform_request(
&self,
physical_region_id: RegionId,
logical_region_id: RegionId,
mut request: ScanRequest,
request: ScanRequest,
) -> Result<ScanRequest> {
// transform projection
let physical_projection = match request.projection_input.as_ref() {
Some(projection_input) => {
self.transform_projection(
physical_region_id,
logical_region_id,
&projection_input.projection,
)
.await?
}
None => {
self.default_projection(physical_region_id, logical_region_id)
.await?
}
};
// Rewrite the top-level projection from logical-region schema indices to
// physical-region schema indices. `nested_paths` are left unchanged because
// they are expressed by column name rather than schema index.
request.projection_input.get_or_insert_default().projection = physical_projection;
request
.filters
.push(self.table_id_filter(logical_region_id));
Ok(request)
let logical_metadata = self
.logical_region_metadata(physical_region_id, logical_region_id)
.await?;
self.transform_logical_request_with_mapper(
physical_region_id,
logical_region_id,
request,
&logical_metadata,
)
.await
.map(|(request, _)| request)
}
/// Generate a filter on the table id column.
fn table_id_filter(&self, logical_region_id: RegionId) -> Expr {
logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
.eq(logical_expr::lit(logical_region_id.table_id()))
}
/// Transform the projection from logical region to physical region.
///
/// This method will not preserve internal columns.
pub async fn transform_projection(
async fn transform_logical_request_with_mapper(
&self,
physical_region_id: RegionId,
logical_region_id: RegionId,
request: ScanRequest,
logical_metadata: &RegionMetadataRef,
) -> Result<(ScanRequest, ValueSplitProjectionMapper)> {
let data_region_id = utils::to_data_region_id(physical_region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
.await
.context(MitoReadOperationSnafu)?;
self.transform_request_with_mapper(
logical_region_id,
request,
logical_metadata,
&physical_metadata,
Some(logical_region_id),
)
}
fn transform_request_with_mapper(
&self,
region_id: RegionId,
mut request: ScanRequest,
visible_metadata: &RegionMetadataRef,
physical_metadata: &RegionMetadataRef,
logical_region_id: Option<RegionId>,
) -> Result<(ScanRequest, ValueSplitProjectionMapper)> {
let split_value_columns = split_value_columns(visible_metadata, physical_metadata);
let mut residual_column_names = HashSet::new();
let residual_filters = request
.filters
.iter()
.filter(|filter| {
let mut columns = HashSet::new();
let is_residual = expr_to_columns(filter, &mut columns).is_ok()
&& columns
.iter()
.any(|column| split_value_columns.contains(&column.name));
if is_residual {
residual_column_names.extend(columns.into_iter().map(|column| column.name));
}
is_residual
})
.cloned()
.collect::<Vec<_>>();
let mut visible_projection = match request.projection_input.as_ref() {
Some(projection_input) => projection_input.projection.clone(),
None => (0..visible_metadata.column_metadatas.len()).collect(),
};
let visible_columns = visible_projection.len();
let mut projected = visible_projection.iter().copied().collect::<HashSet<_>>();
for (index, column) in visible_metadata.column_metadatas.iter().enumerate() {
if residual_column_names.contains(&column.column_schema.name) && projected.insert(index)
{
visible_projection.push(index);
}
}
let (physical_projection, mapper) = self.transform_projection_with_mapper(
region_id,
&visible_projection,
visible_metadata,
physical_metadata,
visible_columns,
residual_filters,
)?;
// Top-level projections are indices; nested paths are column names.
request.projection_input.get_or_insert_default().projection = physical_projection;
request.filters = request
.filters
.into_iter()
.map(|filter| rewrite_metric_value_filter(region_id, filter, &split_value_columns))
.collect::<Result<Vec<_>>>()?;
if let Some(logical_region_id) = logical_region_id {
request.filters.push(
logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
.eq(logical_expr::lit(logical_region_id.table_id())),
);
}
Ok((request, mapper))
}
fn transform_projection_with_mapper(
&self,
logical_region_id: RegionId,
origin_projection: &[usize],
) -> Result<Vec<usize>> {
// project on logical columns
let all_logical_columns = self
.load_logical_column_names(physical_region_id, logical_region_id)
.await?;
let projected_logical_names = origin_projection
.iter()
.map(|i| all_logical_columns[*i].clone())
.collect::<Vec<_>>();
// generate physical projection
logical_metadata: &RegionMetadataRef,
physical_metadata: &RegionMetadataRef,
visible_columns: usize,
residual_filters: Vec<Expr>,
) -> Result<(Vec<usize>, ValueSplitProjectionMapper)> {
let mut physical_projection = Vec::with_capacity(origin_projection.len());
let data_region_id = utils::to_data_region_id(physical_region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
.await
.context(MitoReadOperationSnafu)?;
let mut output_columns = Vec::with_capacity(origin_projection.len());
for name in projected_logical_names {
for logical_idx in origin_projection {
let logical_column = logical_metadata
.column_metadatas
.get(*logical_idx)
.with_context(|| InvalidRequestSnafu {
region_id: logical_region_id,
reason: format!("projection index {} is out of bound", logical_idx),
})?;
let name = &logical_column.column_schema.name;
// Safety: logical columns is a strict subset of physical columns
physical_projection.push(physical_metadata.column_index_by_name(&name).unwrap());
let float_index = physical_metadata.column_index_by_name(name).unwrap();
let input_float_index = physical_projection.len();
physical_projection.push(float_index);
if logical_column.semantic_type == SemanticType::Field
&& logical_column.column_schema.data_type == ConcreteDataType::float64_datatype()
&& let Some(int_index) = physical_metadata
.column_index_by_name(&metric_engine_value_int_column_name(name))
{
let input_int_index = physical_projection.len();
physical_projection.push(int_index);
output_columns.push(ValueColumnProjection::Split {
float_index: input_float_index,
int_index: input_int_index,
output_schema: logical_column.column_schema.clone(),
});
} else {
output_columns.push(ValueColumnProjection::Direct {
input_index: input_float_index,
output_schema: logical_column.column_schema.clone(),
});
}
}
Ok(physical_projection)
}
/// Default projection for a logical region. Includes non-internal columns
pub async fn default_projection(
&self,
physical_region_id: RegionId,
logical_region_id: RegionId,
) -> Result<Vec<usize>> {
let logical_columns = self
.load_logical_column_names(physical_region_id, logical_region_id)
.await?;
let mut projection = Vec::with_capacity(logical_columns.len());
let data_region_id = utils::to_data_region_id(physical_region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
.await
.context(MitoReadOperationSnafu)?;
for name in logical_columns {
// Safety: logical columns is a strict subset of physical columns
projection.push(physical_metadata.column_index_by_name(&name).unwrap());
}
Ok(projection)
Ok((
physical_projection,
ValueSplitProjectionMapper::new(output_columns, visible_columns, residual_filters),
))
}
pub async fn logical_region_metadata(
@@ -269,6 +379,119 @@ impl MetricEngineInner {
}
}
fn visible_physical_region_metadata(
physical_metadata: &RegionMetadataRef,
) -> Result<RegionMetadataRef> {
let visible_columns = physical_metadata
.column_metadatas
.iter()
.filter(|column| !is_metric_engine_value_int_column(&column.column_schema.name))
.cloned()
.collect::<Vec<_>>();
if visible_columns.len() == physical_metadata.column_metadatas.len() {
return Ok(physical_metadata.clone());
}
let primary_key = physical_metadata.primary_key.clone();
let mut builder = RegionMetadataBuilder::from_existing((**physical_metadata).clone());
builder
.alter(AlterKind::SyncColumns {
column_metadatas: visible_columns,
})
.context(InvalidMetadataSnafu)?;
builder.primary_key(primary_key);
builder.build().map(Arc::new).context(InvalidMetadataSnafu)
}
fn split_value_columns(
logical_metadata: &RegionMetadataRef,
physical_metadata: &RegionMetadataRef,
) -> HashSet<String> {
logical_metadata
.column_metadatas
.iter()
.filter(|column| {
column.semantic_type == SemanticType::Field
&& column.column_schema.data_type == ConcreteDataType::float64_datatype()
})
.filter_map(|column| {
let value_name = &column.column_schema.name;
physical_metadata
.column_by_name(&metric_engine_value_int_column_name(value_name))
.filter(|int_column| {
int_column.semantic_type == SemanticType::Field
&& int_column.column_schema.data_type == ConcreteDataType::int64_datatype()
})
.map(|_| value_name.clone())
})
.collect()
}
fn rewrite_metric_value_filter(
logical_region_id: RegionId,
filter: Expr,
split_value_columns: &HashSet<String>,
) -> Result<Expr> {
if split_value_columns.is_empty() {
return Ok(filter);
}
let filter_display = filter.to_string();
let mut rewriter = MetricValueFilterRewriter {
split_value_columns,
};
filter
.rewrite(&mut rewriter)
.map(|rewritten| rewritten.data)
.map_err(|err| {
InvalidRequestSnafu {
region_id: logical_region_id,
reason: format!("failed to rewrite metric value filter {filter_display}: {err}"),
}
.build()
})
}
struct MetricValueFilterRewriter<'a> {
split_value_columns: &'a HashSet<String>,
}
impl TreeNodeRewriter for MetricValueFilterRewriter<'_> {
type Node = Expr;
fn f_down(&mut self, expr: Expr) -> DataFusionResult<Transformed<Expr>> {
let recursion = if matches!(
expr,
Expr::Exists(_) | Expr::InSubquery(_) | Expr::ScalarSubquery(_)
) {
TreeNodeRecursion::Jump
} else {
TreeNodeRecursion::Continue
};
Ok(Transformed::new(expr, false, recursion))
}
fn f_up(&mut self, expr: Expr) -> DataFusionResult<Transformed<Expr>> {
let Expr::Column(column) = expr else {
return Ok(Transformed::no(expr));
};
if !self.split_value_columns.contains(&column.name) {
return Ok(Transformed::no(Expr::Column(column)));
}
let int_column = Column {
relation: column.relation.clone(),
name: metric_engine_value_int_column_name(&column.name),
spans: column.spans.clone(),
};
let float_expr = Expr::Column(column);
let int_expr = cast(Expr::Column(int_column), ArrowDataType::Float64);
Ok(Transformed::yes(coalesce(vec![int_expr, float_expr])))
}
}
#[cfg(test)]
impl MetricEngineInner {
pub async fn scan_to_stream(
@@ -276,27 +499,23 @@ impl MetricEngineInner {
region_id: RegionId,
request: ScanRequest,
) -> Result<common_recordbatch::SendableRecordBatchStream, common_error::ext::BoxedError> {
let is_reading_physical_region = self.is_physical_region(region_id);
if is_reading_physical_region {
self.mito
.scan_to_stream(region_id, request)
.await
.map_err(common_error::ext::BoxedError::new)
} else {
let physical_region_id = self
.get_physical_region_id(region_id)
.await
.map_err(common_error::ext::BoxedError::new)?;
let request = self
.transform_request(physical_region_id, region_id, request)
.await
.map_err(common_error::ext::BoxedError::new)?;
self.mito
.scan_to_stream(physical_region_id, request)
.await
.map_err(common_error::ext::BoxedError::new)
}
let scanner = self
.read_region(region_id, request)
.await
.map_err(common_error::ext::BoxedError::new)?;
let metrics_set = datafusion::physical_plan::metrics::ExecutionPlanMetricsSet::new();
let streams = (0..scanner.properties().num_partitions())
.map(|partition| {
scanner.scan_partition(
&store_api::region_engine::QueryScanContext::default(),
&metrics_set,
partition,
)
})
.collect::<std::result::Result<Vec<_>, _>>()?;
common_recordbatch::util::ChainedRecordBatchStream::new(streams)
.map(|stream| Box::pin(stream) as _)
.map_err(common_error::ext::BoxedError::new)
}
}
@@ -351,7 +570,7 @@ mod test {
assert_eq!(
scan_req.projection_indices().unwrap(),
&[11, 10, 9, 8, 0, 1, 4]
&[12, 11, 10, 9, 0, 1, 2, 5]
);
assert_eq!(scan_req.filters.len(), 1);
assert_eq!(
@@ -370,7 +589,7 @@ mod test {
.unwrap();
assert_eq!(
scan_req.projection_indices().unwrap(),
&[11, 10, 9, 8, 0, 1, 4]
&[12, 11, 10, 9, 0, 1, 2, 5]
);
}
}
@@ -79,37 +79,4 @@ impl MetricEngineInner {
Ok(dedup_columns)
}
/// Load logical column names of a logical region.
///
/// The return value is ordered on column name alphabetically.
pub async fn load_logical_column_names(
&self,
physical_region_id: RegionId,
logical_region_id: RegionId,
) -> Result<Vec<String>> {
// First try to load from state cache
if let Some(columns) = self
.state
.read()
.unwrap()
.logical_columns()
.get(&logical_region_id)
{
return Ok(columns
.iter()
.map(|c| c.column_schema.name.clone())
.collect());
}
// Else load from metadata region
let columns = self
.load_logical_columns(physical_region_id, logical_region_id)
.await?
.into_iter()
.map(|c| c.column_schema.name)
.collect::<Vec<_>>();
Ok(columns)
}
}
+1
View File
@@ -64,3 +64,4 @@ pub mod row_modifier;
#[cfg(test)]
mod test_util;
mod utils;
mod value_split;
+7 -3
View File
@@ -30,6 +30,7 @@ use store_api::metadata::ColumnMetadata;
use store_api::metric_engine_consts::{
ALTER_PHYSICAL_EXTENSION_KEY, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME,
PHYSICAL_TABLE_METADATA_KEY, TABLE_COLUMN_METADATA_EXTENSION_KEY,
metric_engine_value_int_column_name,
};
use store_api::path_utils::table_dir;
use store_api::region_engine::RegionEngine;
@@ -204,7 +205,7 @@ impl TestEnv {
let column_metadatas =
parse_column_metadatas(&response.extensions, TABLE_COLUMN_METADATA_EXTENSION_KEY)
.unwrap();
assert_eq!(column_metadatas.len(), 4);
assert_eq!(column_metadatas.len(), 5);
}
/// Create logical region in [MetricEngine] with specific `physical_region_id` and `logical_region_id`.
@@ -228,7 +229,7 @@ impl TestEnv {
.unwrap();
let column_metadatas =
parse_column_metadatas(&response.extensions, ALTER_PHYSICAL_EXTENSION_KEY).unwrap();
assert_eq!(column_metadatas.len(), 5);
assert_eq!(column_metadatas.len(), 6);
let column_names = column_metadatas
.iter()
.map(|c| c.column_schema.name.as_str())
@@ -237,11 +238,13 @@ impl TestEnv {
.iter()
.map(|c| c.column_id)
.collect::<Vec<_>>();
let value_int_name = metric_engine_value_int_column_name(greptime_value());
assert_eq!(
column_names,
vec![
greptime_timestamp(),
greptime_value(),
value_int_name.as_str(),
"__table_id",
"__tsid",
"job",
@@ -252,9 +255,10 @@ impl TestEnv {
vec![
0,
1,
2,
ReservedColumnId::table_id(),
ReservedColumnId::tsid(),
2,
3,
]
);
}
+325
View File
@@ -0,0 +1,325 @@
// 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::fmt;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use common_error::ext::BoxedError;
use common_recordbatch::adapter::RecordBatchMetrics;
use common_recordbatch::error::{
CreateRecordBatchesSnafu, NewDfRecordBatchSnafu, PhysicalExprSnafu, Result as RecordBatchResult,
};
use common_recordbatch::filter::batch_filter;
use common_recordbatch::{
DfRecordBatch, OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
};
use datafusion::common::ToDFSchema;
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::Expr;
use datafusion::logical_expr::utils::conjunction;
use datafusion::physical_expr::create_physical_expr;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{DisplayAs, DisplayFormatType};
use datafusion_physical_expr::PhysicalExpr;
use datatypes::arrow::array::{Array, ArrayRef, Float64Array, Float64Builder, Int64Array};
use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
use futures_util::Stream;
use snafu::{OptionExt, ResultExt};
use store_api::metadata::RegionMetadataRef;
use store_api::region_engine::{
PrepareRequest, QueryScanContext, RegionScanner, RegionScannerRef, ScannerProperties,
};
use store_api::storage::{RegionId, SequenceNumber};
#[derive(Clone)]
pub(crate) enum ValueColumnProjection {
Direct {
input_index: usize,
output_schema: ColumnSchema,
},
Split {
float_index: usize,
int_index: usize,
output_schema: ColumnSchema,
},
}
#[derive(Clone)]
pub(crate) struct ValueSplitProjectionMapper {
output_schema: SchemaRef,
working_schema: SchemaRef,
columns: Vec<ValueColumnProjection>,
visible_columns: usize,
has_split: bool,
residual_filters: Vec<Expr>,
}
impl ValueSplitProjectionMapper {
pub(crate) fn new(
columns: Vec<ValueColumnProjection>,
visible_columns: usize,
residual_filters: Vec<Expr>,
) -> Self {
let has_split = columns
.iter()
.any(|column| matches!(column, ValueColumnProjection::Split { .. }));
let working_columns = columns
.iter()
.map(|column| match column {
ValueColumnProjection::Direct { output_schema, .. }
| ValueColumnProjection::Split { output_schema, .. } => output_schema.clone(),
})
.collect::<Vec<_>>();
let output_columns = working_columns
.iter()
.take(visible_columns)
.cloned()
.collect::<Vec<_>>();
Self {
output_schema: Arc::new(Schema::new(output_columns)),
working_schema: Arc::new(Schema::new(working_columns)),
columns,
visible_columns,
has_split,
residual_filters,
}
}
fn convert_batch(&self, batch: RecordBatch) -> RecordBatchResult<RecordBatch> {
if !self.has_split && self.residual_filters.is_empty() {
let projection = self
.columns
.iter()
.map(|column| match column {
ValueColumnProjection::Direct { input_index, .. } => *input_index,
ValueColumnProjection::Split { .. } => unreachable!(),
})
.collect::<Vec<_>>();
return batch.try_project(&projection);
}
let arrays = self
.columns
.iter()
.map(|column| match column {
ValueColumnProjection::Direct { input_index, .. } => {
Ok(batch.column(*input_index).clone())
}
ValueColumnProjection::Split {
float_index,
int_index,
..
} => coalesce_value_columns(batch.column(*float_index), batch.column(*int_index)),
})
.collect::<RecordBatchResult<Vec<_>>>()?;
let df_record_batch =
DfRecordBatch::try_new(self.working_schema.arrow_schema().clone(), arrays)
.context(NewDfRecordBatchSnafu)?;
let mut batch =
RecordBatch::from_df_record_batch(self.working_schema.clone(), df_record_batch);
batch = self.apply_residual_filters(batch)?;
if self.visible_columns == self.columns.len() {
Ok(RecordBatch::from_df_record_batch(
self.output_schema.clone(),
batch.into_df_record_batch(),
))
} else {
let projection = (0..self.visible_columns).collect::<Vec<_>>();
batch.try_project(&projection)
}
}
fn apply_residual_filters(&self, batch: RecordBatch) -> RecordBatchResult<RecordBatch> {
let Some(filter) = conjunction(self.residual_filters.clone()) else {
return Ok(batch);
};
let df_schema = self
.working_schema
.arrow_schema()
.clone()
.to_dfschema_ref()
.context(PhysicalExprSnafu)?;
let predicate = create_physical_expr(&filter, &df_schema, &ExecutionProps::new())
.context(PhysicalExprSnafu)?;
let df_record_batch =
batch_filter(batch.df_record_batch(), &predicate).context(PhysicalExprSnafu)?;
Ok(RecordBatch::from_df_record_batch(
self.working_schema.clone(),
df_record_batch,
))
}
}
fn coalesce_value_columns(float_col: &ArrayRef, int_col: &ArrayRef) -> RecordBatchResult<ArrayRef> {
let float_array = float_col
.as_any()
.downcast_ref::<Float64Array>()
.with_context(|| CreateRecordBatchesSnafu {
reason: format!("expected Float64 metric value column, got {float_col:?}"),
})?;
let int_array = int_col
.as_any()
.downcast_ref::<Int64Array>()
.with_context(|| CreateRecordBatchesSnafu {
reason: format!("expected Int64 metric value column, got {int_col:?}"),
})?;
let mut builder = Float64Builder::with_capacity(float_array.len());
for row in 0..float_array.len() {
if !int_array.is_null(row) {
builder.append_value(int_array.value(row) as f64);
} else if !float_array.is_null(row) {
builder.append_value(float_array.value(row));
} else {
builder.append_null();
}
}
Ok(Arc::new(builder.finish()))
}
pub(crate) struct ValueSplitScanner {
inner: RegionScannerRef,
logical_metadata: RegionMetadataRef,
mapper: ValueSplitProjectionMapper,
}
impl ValueSplitScanner {
pub(crate) fn new(
inner: RegionScannerRef,
logical_metadata: RegionMetadataRef,
mapper: ValueSplitProjectionMapper,
) -> Self {
Self {
inner,
logical_metadata,
mapper,
}
}
}
impl fmt::Debug for ValueSplitScanner {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ValueSplitScanner")
.field("inner", &self.inner)
.field("has_split", &self.mapper.has_split)
.finish()
}
}
impl DisplayAs for ValueSplitScanner {
fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
self.inner.fmt_as(t, f)
}
}
impl RegionScanner for ValueSplitScanner {
fn name(&self) -> &str {
self.inner.name()
}
fn properties(&self) -> &ScannerProperties {
self.inner.properties()
}
fn schema(&self) -> SchemaRef {
self.mapper.output_schema.clone()
}
fn metadata(&self) -> RegionMetadataRef {
self.logical_metadata.clone()
}
fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> {
self.inner.prepare(request)
}
fn scan_partition(
&self,
ctx: &QueryScanContext,
metrics_set: &ExecutionPlanMetricsSet,
partition: usize,
) -> Result<SendableRecordBatchStream, BoxedError> {
let stream = self.inner.scan_partition(ctx, metrics_set, partition)?;
Ok(Box::pin(ValueSplitRecordBatchStream {
inner: stream,
mapper: self.mapper.clone(),
}))
}
fn has_predicate_without_region(&self) -> bool {
self.inner.has_predicate_without_region()
}
fn add_dyn_filter_to_predicate(
&mut self,
filter_exprs: Vec<Arc<dyn PhysicalExpr>>,
) -> Vec<bool> {
if self.mapper.has_split {
return vec![false; filter_exprs.len()];
}
self.inner.add_dyn_filter_to_predicate(filter_exprs)
}
fn set_logical_region(&mut self, logical_region: bool) {
self.inner.set_logical_region(logical_region);
}
fn set_query_load_region_id(&mut self, region_id: RegionId) {
self.inner.set_query_load_region_id(region_id);
}
fn snapshot_sequence(&self) -> Option<SequenceNumber> {
self.inner.snapshot_sequence()
}
}
struct ValueSplitRecordBatchStream {
inner: SendableRecordBatchStream,
mapper: ValueSplitProjectionMapper,
}
impl RecordBatchStream for ValueSplitRecordBatchStream {
fn name(&self) -> &str {
"ValueSplitRecordBatchStream"
}
fn schema(&self) -> SchemaRef {
self.mapper.output_schema.clone()
}
fn output_ordering(&self) -> Option<&[OrderOption]> {
self.inner.output_ordering()
}
fn metrics(&self) -> Option<RecordBatchMetrics> {
self.inner.metrics()
}
}
impl Stream for ValueSplitRecordBatchStream {
type Item = RecordBatchResult<RecordBatch>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.inner)
.poll_next(cx)
.map(|opt| opt.map(|result| result.and_then(|batch| self.mapper.convert_batch(batch))))
}
}
+31 -2
View File
@@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex};
use common_query::stream::StreamScanAdapter;
use common_recordbatch::OrderOption;
use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
use datafusion::arrow::datatypes::{Schema as DfSchema, SchemaRef as DfSchemaRef};
use datafusion::catalog::Session;
use datafusion::datasource::{TableProvider, TableType as DfTableType};
use datafusion::error::Result as DfResult;
@@ -27,6 +27,9 @@ use datafusion_expr::TableProviderFilterPushDown as DfTableProviderFilterPushDow
use datafusion_expr::expr::Expr;
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_expr::expressions::Column;
use store_api::metric_engine_consts::{
METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, is_metric_engine_value_int_column,
};
use store_api::storage::{ScanRequest, VectorSearchRequest};
use crate::table::{TableRef, TableType};
@@ -82,7 +85,33 @@ impl TableProvider for DfTableProviderAdapter {
}
fn schema(&self) -> DfSchemaRef {
self.table.schema().arrow_schema().clone()
let table_schema = self.table.schema();
let schema = table_schema.arrow_schema();
let table_info = self.table.table_info();
if table_info.meta.engine != METRIC_ENGINE_NAME
|| !table_info
.meta
.options
.extra_options
.contains_key(PHYSICAL_TABLE_METADATA_KEY)
|| !schema
.fields()
.iter()
.any(|field| is_metric_engine_value_int_column(field.name()))
{
return schema.clone();
}
let fields = schema
.fields()
.iter()
.filter(|field| !is_metric_engine_value_int_column(field.name()))
.cloned()
.collect::<Vec<_>>();
Arc::new(DfSchema::new_with_metadata(
fields,
schema.metadata().clone(),
))
}
fn table_type(&self) -> DfTableType {
@@ -50,16 +50,17 @@ DESC TABLE t2;
DESC TABLE phy;
+------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+------------+----------------------+-----+------+---------+---------------+
+-----------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+-----------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| val__metric_int | Int64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+-----------------+----------------------+-----+------+---------+---------------+
ALTER TABLE t1 ADD COLUMN k STRING PRIMARY KEY;
@@ -93,17 +94,18 @@ DESC TABLE t2;
DESC TABLE phy;
+------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
| k | String | PRI | YES | | TAG |
+------------+----------------------+-----+------+---------+---------------+
+-----------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+-----------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| val__metric_int | Int64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
| k | String | PRI | YES | | TAG |
+-----------------+----------------------+-----+------+---------+---------------+
DROP TABLE t1;
@@ -65,16 +65,17 @@ SELECT table_catalog, table_schema, table_name, table_type, engine FROM informat
DESC TABLE phy;
+------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+------------+----------------------+-----+------+---------+---------------+
+-----------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+-----------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| val__metric_int | Int64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+-----------------+----------------------+-----+------+---------+---------------+
SHOW CREATE TABLE phy;
@@ -125,16 +126,17 @@ Error: 1004(InvalidArguments), Physical region is busy, there are still some log
-- metadata should be restored
DESC TABLE phy;
+------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+------------+----------------------+-----+------+---------+---------------+
+-----------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+-----------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| val__metric_int | Int64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+-----------------+----------------------+-----+------+---------+---------------+
DROP TABLE t1;
@@ -51,16 +51,17 @@ Affected Rows: 0
DESC TABLE phy;
+------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+------------+----------------------+-----+------+---------+---------------+
+-----------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+-----------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| val__metric_int | Int64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+-----------------+----------------------+-----+------+---------+---------------+
SELECT ts, val, __tsid, host, job FROM phy;
@@ -204,16 +205,17 @@ Affected Rows: 0
DESC TABLE phy;
+------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+------------+----------------------+-----+------+---------+---------------+
+-----------------+----------------------+-----+------+---------+---------------+
| Column | Type | Key | Null | Default | Semantic Type |
+-----------------+----------------------+-----+------+---------+---------------+
| ts | TimestampMillisecond | PRI | NO | | TIMESTAMP |
| val | Float64 | | YES | | FIELD |
| val__metric_int | Int64 | | YES | | FIELD |
| __table_id | UInt32 | PRI | NO | | TAG |
| __tsid | UInt64 | PRI | NO | | TAG |
| host | String | PRI | YES | | TAG |
| job | String | PRI | YES | | TAG |
+-----------------+----------------------+-----+------+---------+---------------+
DROP TABLE phy;