refactor(mito2): own split metric value reads

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2026-07-13 20:16:39 +08:00
parent f8a2d3f23b
commit c79581ecf8
15 changed files with 495 additions and 347 deletions
Generated
-1
View File
@@ -8332,7 +8332,6 @@ dependencies = [
"common-wal",
"criterion 0.7.0",
"datafusion",
"datafusion-physical-expr",
"datatypes",
"futures-util",
"fxhash",
-1
View File
@@ -26,7 +26,6 @@ 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 -4
View File
@@ -101,7 +101,7 @@ impl DataRegion {
// retrieve underlying version
let region_metadata = self
.mito
.get_metadata(region_id)
.get_physical_metadata(region_id)
.await
.context(MitoReadOperationSnafu)?;
@@ -214,7 +214,7 @@ impl DataRegion {
let data_region_id = utils::to_data_region_id(physical_region_id);
let metadata = self
.mito
.get_metadata(data_region_id)
.get_physical_metadata(data_region_id)
.await
.context(MitoReadOperationSnafu)?;
Ok(metadata.column_metadatas.clone())
@@ -268,7 +268,7 @@ mod test {
let current_version = env
.mito()
.get_metadata(utils::to_data_region_id(env.default_physical_region_id()))
.get_physical_metadata(utils::to_data_region_id(env.default_physical_region_id()))
.await
.unwrap()
.schema_version;
@@ -306,7 +306,7 @@ mod test {
let new_metadata = env
.mito()
.get_metadata(utils::to_data_region_id(env.default_physical_region_id()))
.get_physical_metadata(utils::to_data_region_id(env.default_physical_region_id()))
.await
.unwrap();
let column_names = new_metadata
+2 -3
View File
@@ -25,7 +25,6 @@ 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;
@@ -163,7 +162,7 @@ impl MetricEngineInner {
) -> Result<RecordBatch> {
let metadata = self
.mito
.get_metadata(data_region_id)
.get_physical_metadata(data_region_id)
.await
.context(error::MitoReadOperationSnafu)?;
let batch_schema = batch.schema();
@@ -222,7 +221,7 @@ impl MetricEngineInner {
async fn physical_schema_version(&self, region_id: RegionId) -> Result<u64> {
Ok(self
.mito
.get_metadata(region_id)
.get_physical_metadata(region_id)
.await
.context(error::MitoReadOperationSnafu)?
.schema_version)
+33 -1
View File
@@ -1251,6 +1251,7 @@ mod tests {
.unwrap();
let physical_region_id = env.default_physical_region_id();
let data_region_id = to_data_region_id(physical_region_id);
let visible_physical_stream = env
.metric()
.scan_to_stream(physical_region_id, ScanRequest::default())
@@ -1258,6 +1259,24 @@ mod tests {
.unwrap();
let int_column_name =
store_api::metric_engine_consts::metric_engine_value_int_column_name(greptime_value());
let visible_mito_metadata = env.mito().get_metadata(data_region_id).await.unwrap();
let physical_mito_metadata = env
.mito()
.get_physical_metadata(data_region_id)
.await
.unwrap();
assert!(
visible_mito_metadata
.column_by_name(&int_column_name)
.is_none(),
"mito query metadata should hide split companion columns"
);
assert!(
physical_mito_metadata
.column_by_name(&int_column_name)
.is_some(),
"mito physical metadata should retain split companion columns"
);
assert!(
visible_physical_stream
.schema()
@@ -1266,6 +1285,20 @@ mod tests {
"metric physical reads should hide split companion column"
);
let logical_scanner = env
.metric()
.handle_query(logical_region_id, ScanRequest::default())
.await
.unwrap();
assert_eq!(logical_scanner.metadata().region_id, logical_region_id);
assert!(
logical_scanner
.metadata()
.column_by_name(&int_column_name)
.is_none(),
"logical scanner metadata should not expose split companion columns"
);
let logical_batches = RecordBatches::try_collect(
env.metric()
.scan_to_stream(logical_region_id, ScanRequest::default())
@@ -1352,7 +1385,6 @@ mod tests {
("integer".to_string(), 2, None, Some(3)),
("integer".to_string(), 3, None, Some(4)),
];
let data_region_id = to_data_region_id(physical_region_id);
let raw_physical_batches = RecordBatches::try_collect(
env.mito()
.scan_to_stream(data_region_id, ScanRequest::default())
+121 -277
View File
@@ -12,27 +12,24 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::fmt;
use std::sync::Arc;
use api::v1::SemanticType;
use common_error::ext::BoxedError;
use common_recordbatch::SendableRecordBatchStream;
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 datafusion::logical_expr;
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{DisplayAs, DisplayFormatType};
use snafu::{OptionExt, ResultExt};
use store_api::metadata::{RegionMetadataBuilder, RegionMetadataRef};
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::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME;
use store_api::region_engine::{
PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, RegionScannerRef,
ScannerProperties,
};
use store_api::region_engine::{RegionEngine, RegionScannerRef};
use store_api::region_request::AlterKind;
use store_api::storage::{RegionId, ScanRequest, SequenceNumber};
use crate::engine::MetricEngineInner;
@@ -42,7 +39,6 @@ use crate::error::{
};
use crate::metrics::MITO_OPERATION_ELAPSED;
use crate::utils;
use crate::value_split::{ValueColumnProjection, ValueSplitProjectionMapper, ValueSplitScanner};
impl MetricEngineInner {
#[tracing::instrument(skip_all)]
@@ -73,32 +69,10 @@ impl MetricEngineInner {
.with_label_values(&["read_physical"])
.start_timer();
let data_region_id = utils::to_data_region_id(region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
self.mito
.handle_query(region_id, request)
.await
.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,
)))
.context(MitoReadOperationSnafu)
}
async fn read_logical_region(
@@ -115,8 +89,8 @@ impl MetricEngineInner {
let logical_metadata = self
.logical_region_metadata(physical_region_id, logical_region_id)
.await?;
let (request, mapper) = self
.transform_logical_request_with_mapper(
let request = self
.transform_request_with_metadata(
physical_region_id,
logical_region_id,
request,
@@ -131,10 +105,9 @@ impl MetricEngineInner {
scanner.set_logical_region(true);
scanner.set_query_load_region_id(data_region_id);
Ok(Box::new(ValueSplitScanner::new(
Ok(Box::new(LogicalRegionScanner::new(
scanner,
logical_metadata,
mapper,
)))
}
@@ -156,13 +129,10 @@ impl MetricEngineInner {
self.state.read().unwrap().exist_physical_region(region_id);
if is_reading_physical_region {
let data_region_id = utils::to_data_region_id(region_id);
let physical_metadata = self
.mito
.get_metadata(data_region_id)
self.mito
.get_metadata(region_id)
.await
.context(MitoReadOperationSnafu)?;
visible_physical_region_metadata(&physical_metadata)
.context(MitoReadOperationSnafu)
} else {
let physical_region_id = self.get_physical_region_id(region_id).await?;
self.logical_region_metadata(physical_region_id, region_id)
@@ -198,152 +168,60 @@ impl MetricEngineInner {
let logical_metadata = self
.logical_region_metadata(physical_region_id, logical_region_id)
.await?;
self.transform_logical_request_with_mapper(
self.transform_request_with_metadata(
physical_region_id,
logical_region_id,
request,
&logical_metadata,
)
.await
.map(|(request, _)| request)
}
async fn transform_logical_request_with_mapper(
async fn transform_request_with_metadata(
&self,
physical_region_id: RegionId,
logical_region_id: RegionId,
request: ScanRequest,
mut request: ScanRequest,
logical_metadata: &RegionMetadataRef,
) -> Result<(ScanRequest, ValueSplitProjectionMapper)> {
) -> Result<ScanRequest> {
let logical_projection = match request.projection_input.as_ref() {
Some(projection_input) => projection_input.projection.clone(),
None => (0..logical_metadata.column_metadatas.len()).collect(),
};
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 mut physical_projection = Vec::with_capacity(logical_projection.len());
for logical_index in logical_projection {
let logical_column = logical_metadata
.column_metadatas
.get(logical_index)
.with_context(|| InvalidRequestSnafu {
region_id: logical_region_id,
reason: format!("projection index {logical_index} is out of bound"),
})?;
let name = &logical_column.column_schema.name;
let physical_index =
physical_metadata
.column_index_by_name(name)
.with_context(|| InvalidRequestSnafu {
region_id: logical_region_id,
reason: format!("column {name} is missing from physical metadata"),
})?;
physical_projection.push(physical_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],
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 mut output_columns = Vec::with_capacity(origin_projection.len());
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
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,
ValueSplitProjectionMapper::new(output_columns, visible_columns, residual_filters),
))
request.filters.push(
logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
.eq(logical_expr::lit(logical_region_id.table_id())),
);
Ok(request)
}
pub async fn logical_region_metadata(
@@ -379,116 +257,82 @@ 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());
struct LogicalRegionScanner {
inner: RegionScannerRef,
metadata: RegionMetadataRef,
}
impl LogicalRegionScanner {
fn new(inner: RegionScannerRef, metadata: RegionMetadataRef) -> Self {
Self { inner, metadata }
}
}
impl fmt::Debug for LogicalRegionScanner {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("LogicalRegionScanner")
.field("inner", &self.inner)
.finish()
}
}
impl DisplayAs for LogicalRegionScanner {
fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
self.inner.fmt_as(t, f)
}
}
impl RegionScanner for LogicalRegionScanner {
fn name(&self) -> &str {
self.inner.name()
}
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);
fn properties(&self) -> &ScannerProperties {
self.inner.properties()
}
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 schema(&self) -> datatypes::schema::SchemaRef {
self.inner.schema()
}
fn f_up(&mut self, expr: Expr) -> DataFusionResult<Transformed<Expr>> {
let Expr::Column(column) = expr else {
return Ok(Transformed::no(expr));
};
fn metadata(&self) -> RegionMetadataRef {
self.metadata.clone()
}
if !self.split_value_columns.contains(&column.name) {
return Ok(Transformed::no(Expr::Column(column)));
}
fn prepare(&mut self, request: PrepareRequest) -> std::result::Result<(), BoxedError> {
self.inner.prepare(request)
}
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])))
fn scan_partition(
&self,
ctx: &QueryScanContext,
metrics_set: &ExecutionPlanMetricsSet,
partition: usize,
) -> std::result::Result<SendableRecordBatchStream, BoxedError> {
self.inner.scan_partition(ctx, metrics_set, partition)
}
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> {
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()
}
}
@@ -570,7 +414,7 @@ mod test {
assert_eq!(
scan_req.projection_indices().unwrap(),
&[12, 11, 10, 9, 0, 1, 2, 5]
&[11, 10, 9, 8, 0, 1, 4]
);
assert_eq!(scan_req.filters.len(), 1);
assert_eq!(
@@ -589,7 +433,7 @@ mod test {
.unwrap();
assert_eq!(
scan_req.projection_indices().unwrap(),
&[12, 11, 10, 9, 0, 1, 2, 5]
&[11, 10, 9, 8, 0, 1, 4]
);
}
}
-1
View File
@@ -64,4 +64,3 @@ pub mod row_modifier;
#[cfg(test)]
mod test_util;
mod utils;
mod value_split;
+1
View File
@@ -22,6 +22,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`.
| `request` | `src/mito2/src/request.rs` | `WriteRequest`/`RegionRequest` types and result channels |
| `wal` | `src/mito2/src/wal.rs` | Write-ahead log wrapper over `log-store` |
| `memtable` | `src/mito2/src/memtable/` | In-memory write buffers (time-series / bulk / partition) |
| `metric_value` | `src/mito2/src/metric_value.rs` | Metric float/int companion discovery and logical schema visibility |
| `flush` | `src/mito2/src/flush.rs` | `FlushScheduler`, `WriteBufferManager`, memtable → SST |
| `compaction` | `src/mito2/src/compaction/` | TWCS picker, strict-window manual picker, compactor, memory control |
| `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store |
+15 -2
View File
@@ -142,6 +142,7 @@ use crate::extension::BoxedExtensionRangeProviderFactory;
use crate::gc::GcLimiterRef;
use crate::manifest::action::RegionEdit;
use crate::memtable::MemtableStats;
use crate::metric_value::visible_region_metadata;
use crate::metrics::{
HANDLE_REQUEST_ELAPSED, SCAN_MEMORY_EXHAUSTED_TOTAL, SCAN_MEMORY_USAGE_BYTES,
SCAN_REQUESTS_REJECTED_TOTAL,
@@ -401,6 +402,14 @@ impl MitoEngine {
.map(|r| r.primary_key_encoding())
}
/// Returns persisted region metadata, including engine-internal columns.
pub async fn get_physical_metadata(
&self,
region_id: RegionId,
) -> std::result::Result<RegionMetadataRef, BoxedError> {
self.inner.get_metadata(region_id).map_err(BoxedError::new)
}
/// Handle substrait query and return a stream of record batches
///
/// Notice that the output stream's ordering is not guranateed. If order
@@ -1258,7 +1267,7 @@ impl RegionEngine for MitoEngine {
) -> Result<RegionScannerRef, BoxedError> {
self.scan_region(region_id, request)
.map_err(BoxedError::new)?
.region_scanner()
.query_region_scanner()
.await
.map_err(BoxedError::new)
}
@@ -1281,7 +1290,11 @@ impl RegionEngine for MitoEngine {
&self,
region_id: RegionId,
) -> std::result::Result<RegionMetadataRef, BoxedError> {
self.inner.get_metadata(region_id).map_err(BoxedError::new)
let metadata = self
.inner
.get_metadata(region_id)
.map_err(BoxedError::new)?;
visible_region_metadata(&metadata).map_err(BoxedError::new)
}
/// Stop the engine.
+1
View File
@@ -36,6 +36,7 @@ pub mod flush;
pub mod gc;
pub mod manifest;
pub mod memtable;
mod metric_value;
mod metrics;
pub mod read;
pub mod region;
+104
View File
@@ -0,0 +1,104 @@
// 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::HashSet;
use std::sync::Arc;
use api::v1::SemanticType;
use datatypes::prelude::ConcreteDataType;
use snafu::ResultExt;
use store_api::metadata::{RegionMetadata, RegionMetadataBuilder, RegionMetadataRef};
use store_api::metric_engine_consts::{
DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
metric_engine_value_int_column_name,
};
use store_api::region_request::AlterKind;
use store_api::storage::consts::ReservedColumnId;
use crate::error::{InvalidMetadataSnafu, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct MetricValueColumn {
pub(crate) value_index: usize,
pub(crate) int_index: usize,
}
pub(crate) fn metric_value_columns(metadata: &RegionMetadata) -> Vec<MetricValueColumn> {
if !is_metric_engine_data_region(metadata) {
return Vec::new();
}
metadata
.field_columns()
.filter(|column| column.column_schema.data_type == ConcreteDataType::float64_datatype())
.filter_map(|value_column| {
let int_name = metric_engine_value_int_column_name(&value_column.column_schema.name);
let int_column = metadata.column_by_name(&int_name)?;
if int_column.semantic_type != SemanticType::Field
|| int_column.column_schema.data_type != ConcreteDataType::int64_datatype()
{
return None;
}
Some(MetricValueColumn {
value_index: metadata.column_index_by_name(&value_column.column_schema.name)?,
int_index: metadata.column_index_by_name(&int_name)?,
})
})
.collect()
}
pub(crate) fn visible_region_metadata(metadata: &RegionMetadataRef) -> Result<RegionMetadataRef> {
let split_columns = metric_value_columns(metadata);
if split_columns.is_empty() {
return Ok(metadata.clone());
}
let int_indices = split_columns
.into_iter()
.map(|column| column.int_index)
.collect::<HashSet<_>>();
let visible_columns = metadata
.column_metadatas
.iter()
.enumerate()
.filter(|(index, _)| !int_indices.contains(index))
.map(|(_, column)| column.clone())
.collect::<Vec<_>>();
let primary_key = metadata.primary_key.clone();
let mut builder = RegionMetadataBuilder::from_existing((**metadata).clone());
builder
.alter(AlterKind::SyncColumns {
column_metadatas: visible_columns,
})
.context(InvalidMetadataSnafu)?;
builder.primary_key(primary_key);
builder.build().map(Arc::new).context(InvalidMetadataSnafu)
}
pub(crate) fn is_metric_engine_data_region(metadata: &RegionMetadata) -> bool {
let has_internal_tag = |name, column_id| {
metadata.column_by_name(name).is_some_and(|column| {
column.semantic_type == SemanticType::Tag
&& column.column_id == column_id
&& metadata.primary_key.contains(&column_id)
})
};
has_internal_tag(
DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
ReservedColumnId::table_id(),
) && has_internal_tag(DATA_SCHEMA_TSID_COLUMN_NAME, ReservedColumnId::tsid())
}
+1
View File
@@ -36,6 +36,7 @@ pub(crate) mod seq_scan;
pub mod series_scan;
pub mod stream;
pub(crate) mod unordered_scan;
pub(crate) mod value_split;
use std::collections::HashMap;
use std::sync::Arc;
+12
View File
@@ -72,6 +72,7 @@ use crate::read::seq_scan::SeqScan;
use crate::read::series_scan::SeriesScan;
use crate::read::stream::ScanBatchStream;
use crate::read::unordered_scan::UnorderedScan;
use crate::read::value_split::{ValueSplitScanner, prepare_value_split_scan};
use crate::read::{BoxedRecordBatchStream, RecordBatch};
use crate::region::options::MergeMode;
use crate::region::version::VersionRef;
@@ -373,6 +374,17 @@ impl ScanRegion {
}
}
/// Returns a public query scanner that hides metric value companion columns.
pub(crate) async fn query_region_scanner(mut self) -> Result<RegionScannerRef> {
let mapper =
prepare_value_split_scan(self.region_id(), &mut self.request, &self.version.metadata)?;
let scanner = self.region_scanner().await?;
match mapper {
Some(mapper) => Ok(Box::new(ValueSplitScanner::new(scanner, mapper))),
None => Ok(scanner),
}
}
/// Scan sequentially.
#[tracing::instrument(skip_all, fields(region_id = %self.region_id()))]
pub(crate) async fn seq_scan(self) -> Result<SeqScan> {
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::fmt;
use std::pin::Pin;
use std::sync::Arc;
@@ -26,23 +27,200 @@ use common_recordbatch::filter::batch_filter;
use common_recordbatch::{
DfRecordBatch, OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
};
use datafusion::common::ToDFSchema;
use datafusion::arrow::datatypes::DataType as ArrowDataType;
use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter};
use datafusion::common::{Column, Result as DataFusionResult, ToDFSchema};
use datafusion::execution::context::ExecutionProps;
use datafusion::functions::expr_fn::coalesce;
use datafusion::logical_expr::Expr;
use datafusion::logical_expr::utils::conjunction;
use datafusion::physical_expr::create_physical_expr;
use datafusion::logical_expr::expr_fn::cast;
use datafusion::logical_expr::utils::{conjunction, expr_to_columns};
use datafusion::physical_expr::{PhysicalExpr, 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 futures::Stream;
use snafu::{OptionExt, ResultExt};
use store_api::metadata::RegionMetadataRef;
use store_api::metric_engine_consts::metric_engine_value_int_column_name;
use store_api::region_engine::{
PrepareRequest, QueryScanContext, RegionScanner, RegionScannerRef, ScannerProperties,
};
use store_api::storage::{RegionId, SequenceNumber};
use store_api::storage::{RegionId, ScanRequest, SequenceNumber};
use crate::error::{InvalidRequestSnafu, Result};
use crate::metric_value::{metric_value_columns, visible_region_metadata};
pub(crate) fn prepare_value_split_scan(
region_id: RegionId,
request: &mut ScanRequest,
physical_metadata: &RegionMetadataRef,
) -> Result<Option<ValueSplitProjectionMapper>> {
let split_columns = metric_value_columns(physical_metadata);
if split_columns.is_empty() {
return Ok(None);
}
let visible_metadata = visible_region_metadata(physical_metadata)?;
let split_value_columns = split_columns
.into_iter()
.map(|column| {
physical_metadata.column_metadatas[column.value_index]
.column_schema
.name
.clone()
})
.collect::<HashSet<_>>();
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 mut physical_projection = Vec::with_capacity(visible_projection.len());
let mut output_columns = Vec::with_capacity(visible_projection.len());
for visible_index in visible_projection {
let visible_column = visible_metadata
.column_metadatas
.get(visible_index)
.with_context(|| InvalidRequestSnafu {
region_id,
reason: format!("projection index {visible_index} is out of bound"),
})?;
let value_name = &visible_column.column_schema.name;
let physical_index = physical_metadata
.column_index_by_name(value_name)
.with_context(|| InvalidRequestSnafu {
region_id,
reason: format!("column {value_name} is missing from physical metadata"),
})?;
let input_value_index = physical_projection.len();
physical_projection.push(physical_index);
if split_value_columns.contains(value_name) {
let int_name = metric_engine_value_int_column_name(value_name);
let int_index = physical_metadata
.column_index_by_name(&int_name)
.with_context(|| InvalidRequestSnafu {
region_id,
reason: format!("column {int_name} is missing from physical metadata"),
})?;
let input_int_index = physical_projection.len();
physical_projection.push(int_index);
output_columns.push(ValueColumnProjection::Split {
float_index: input_value_index,
int_index: input_int_index,
output_schema: visible_column.column_schema.clone(),
});
} else {
output_columns.push(ValueColumnProjection::Direct {
input_index: input_value_index,
output_schema: visible_column.column_schema.clone(),
});
}
}
request.projection_input.get_or_insert_default().projection = physical_projection;
request.filters = request
.filters
.drain(..)
.map(|filter| rewrite_metric_value_filter(region_id, filter, &split_value_columns))
.collect::<Result<Vec<_>>>()?;
Ok(Some(ValueSplitProjectionMapper::new(
visible_metadata,
output_columns,
visible_columns,
residual_filters,
)))
}
fn rewrite_metric_value_filter(
region_id: RegionId,
filter: Expr,
split_value_columns: &HashSet<String>,
) -> Result<Expr> {
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,
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])))
}
}
#[derive(Clone)]
pub(crate) enum ValueColumnProjection {
@@ -59,6 +237,7 @@ pub(crate) enum ValueColumnProjection {
#[derive(Clone)]
pub(crate) struct ValueSplitProjectionMapper {
metadata: RegionMetadataRef,
output_schema: SchemaRef,
working_schema: SchemaRef,
columns: Vec<ValueColumnProjection>,
@@ -69,6 +248,7 @@ pub(crate) struct ValueSplitProjectionMapper {
impl ValueSplitProjectionMapper {
pub(crate) fn new(
metadata: RegionMetadataRef,
columns: Vec<ValueColumnProjection>,
visible_columns: usize,
residual_filters: Vec<Expr>,
@@ -90,6 +270,7 @@ impl ValueSplitProjectionMapper {
.collect::<Vec<_>>();
Self {
metadata,
output_schema: Arc::new(Schema::new(output_columns)),
working_schema: Arc::new(Schema::new(working_columns)),
columns,
@@ -196,21 +377,12 @@ fn coalesce_value_columns(float_col: &ArrayRef, int_col: &ArrayRef) -> RecordBat
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,
}
pub(crate) fn new(inner: RegionScannerRef, mapper: ValueSplitProjectionMapper) -> Self {
Self { inner, mapper }
}
}
@@ -243,7 +415,7 @@ impl RegionScanner for ValueSplitScanner {
}
fn metadata(&self) -> RegionMetadataRef {
self.logical_metadata.clone()
self.mapper.metadata.clone()
}
fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> {
+11 -39
View File
@@ -14,23 +14,17 @@
use std::sync::Arc;
use api::v1::SemanticType;
use datatypes::arrow::array::{
Array, ArrayRef, Float64Array, Float64Builder, Int64Array, Int64Builder,
};
use datatypes::arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer};
use datatypes::arrow::datatypes::SchemaRef;
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::prelude::ConcreteDataType;
use snafu::{OptionExt, ResultExt};
use store_api::metadata::RegionMetadata;
use store_api::metric_engine_consts::{
DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
metric_engine_value_int_column_name,
};
use store_api::storage::consts::ReservedColumnId;
use crate::error::{InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result};
use crate::metric_value::metric_value_columns;
#[derive(Debug, Clone)]
pub(crate) struct MetricValueSplitColumn {
@@ -42,26 +36,15 @@ pub(crate) fn metric_value_split_columns(
metadata: &RegionMetadata,
arrow_schema: &SchemaRef,
) -> Vec<MetricValueSplitColumn> {
if !is_metric_engine_data_region(metadata) {
return vec![];
}
metadata
.field_columns()
.filter(|column| column.column_schema.data_type == ConcreteDataType::float64_datatype())
.filter_map(|float_column| {
let int_name = metric_engine_value_int_column_name(&float_column.column_schema.name);
let int_column = metadata.column_by_name(&int_name)?;
if int_column.semantic_type != SemanticType::Field
|| int_column.column_schema.data_type != ConcreteDataType::int64_datatype()
{
return None;
}
metric_value_columns(metadata)
.into_iter()
.filter_map(|column| {
let float_column = &metadata.column_metadatas[column.value_index];
let int_column = &metadata.column_metadatas[column.int_index];
let float_index = arrow_schema
.index_of(&float_column.column_schema.name)
.ok()?;
let int_index = arrow_schema.index_of(&int_name).ok()?;
let int_index = arrow_schema.index_of(&int_column.column_schema.name).ok()?;
Some(MetricValueSplitColumn {
float_index,
int_index,
@@ -70,21 +53,6 @@ pub(crate) fn metric_value_split_columns(
.collect()
}
fn is_metric_engine_data_region(metadata: &RegionMetadata) -> bool {
let has_internal_tag = |name, column_id| {
metadata.column_by_name(name).is_some_and(|column| {
column.semantic_type == SemanticType::Tag
&& column.column_id == column_id
&& metadata.primary_key.contains(&column_id)
})
};
has_internal_tag(
DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
ReservedColumnId::table_id(),
) && has_internal_tag(DATA_SCHEMA_TSID_COLUMN_NAME, ReservedColumnId::tsid())
}
pub(crate) fn split_metric_value_columns(
batch: &RecordBatch,
split_columns: &[MetricValueSplitColumn],
@@ -247,12 +215,16 @@ fn integer_value(value: f64) -> Option<i64> {
#[cfg(test)]
mod tests {
use api::v1::SemanticType;
use datatypes::arrow::array::{Float64Array, Int64Array};
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::ColumnSchema;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
use store_api::metric_engine_consts::metric_engine_value_int_column_name;
use store_api::storage::RegionId;
use super::*;
use crate::metric_value::is_metric_engine_data_region;
fn column_metadata(
column_id: u32,