metric_engine/engine/

read.rs

// 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::sync::Arc;

use api::v1::SemanticType;
use common_telemetry::{debug, error, tracing};
use datafusion::common::ScalarValue;
use datafusion::logical_expr::{self, Expr, Operator};
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::region_engine::{RegionEngine, RegionScannerRef};
use store_api::storage::{RegionId, ScanRequest, SequenceNumber};

use crate::engine::MetricEngineInner;
use crate::error::{
    InvalidMetadataSnafu, InvalidRequestSnafu, LogicalRegionNotFoundSnafu, MitoReadOperationSnafu,
    Result,
};
use crate::metrics::MITO_OPERATION_ELAPSED;
use crate::utils;

impl MetricEngineInner {
    #[tracing::instrument(skip_all)]
    pub async fn read_region(
        &self,
        region_id: RegionId,
        request: ScanRequest,
    ) -> Result<RegionScannerRef> {
        let is_reading_physical_region = self.is_physical_region(region_id);

        if is_reading_physical_region {
            debug!(
                "Metric region received read request {request:?} on physical region {region_id:?}"
            );
            self.read_physical_region(region_id, request).await
        } else {
            self.read_logical_region(region_id, request).await
        }
    }

    /// Proxy the read request to underlying physical region (mito engine).
    async fn read_physical_region(
        &self,
        region_id: RegionId,
        request: ScanRequest,
    ) -> Result<RegionScannerRef> {
        let _timer = MITO_OPERATION_ELAPSED
            .with_label_values(&["read_physical"])
            .start_timer();

        // PromQL scans physical regions with a table-ID filter to access `__tsid`.
        let logical_region = request
            .filters
            .iter()
            .flat_map(logical_expr::utils::split_conjunction)
            .any(|expr| {
                let Expr::BinaryExpr(expr) = expr else {
                    return false;
                };
                expr.op == Operator::Eq
                    && matches!(
                        (expr.left.as_ref(), expr.right.as_ref()),
                        (Expr::Column(column), Expr::Literal(ScalarValue::UInt32(Some(_)), _))
                            | (Expr::Literal(ScalarValue::UInt32(Some(_)), _), Expr::Column(column))
                            if column.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME
                    )
            });
        let mut scanner = self
            .mito
            .handle_query(region_id, request)
            .await
            .context(MitoReadOperationSnafu)?;
        scanner.set_logical_region(logical_region);

        Ok(scanner)
    }

    async fn read_logical_region(
        &self,
        logical_region_id: RegionId,
        request: ScanRequest,
    ) -> Result<RegionScannerRef> {
        let _timer = MITO_OPERATION_ELAPSED
            .with_label_values(&["read"])
            .start_timer();

        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)
            .await?;
        let mut scanner = self
            .mito
            .handle_query(data_region_id, request)
            .await
            .context(MitoReadOperationSnafu)?;
        scanner.set_logical_region(true);
        scanner.set_query_load_region_id(data_region_id);

        Ok(scanner)
    }

    pub async fn get_last_seq_num(&self, region_id: RegionId) -> Result<SequenceNumber> {
        let region_id = if self.is_physical_region(region_id) {
            region_id
        } else {
            let physical_region_id = self.get_physical_region_id(region_id).await?;
            utils::to_data_region_id(physical_region_id)
        };
        self.mito
            .get_committed_sequence(region_id)
            .await
            .context(MitoReadOperationSnafu)
    }

    pub async fn load_region_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef> {
        let is_reading_physical_region =
            self.state.read().unwrap().exist_physical_region(region_id);

        if is_reading_physical_region {
            self.mito
                .get_metadata(region_id)
                .await
                .context(MitoReadOperationSnafu)
        } else {
            let physical_region_id = self.get_physical_region_id(region_id).await?;
            self.logical_region_metadata(physical_region_id, region_id)
                .await
        }
    }

    /// Returns true if it's a physical region.
    pub fn is_physical_region(&self, region_id: RegionId) -> bool {
        self.state.read().unwrap().exist_physical_region(region_id)
    }

    async fn get_physical_region_id(&self, logical_region_id: RegionId) -> Result<RegionId> {
        let state = &self.state.read().unwrap();
        state
            .get_physical_region_id(logical_region_id)
            .with_context(|| {
                error!("Trying to read an nonexistent region {logical_region_id}");
                LogicalRegionNotFoundSnafu {
                    region_id: logical_region_id,
                }
            })
    }

    /// Transform the [ScanRequest] from logical region to physical data region.
    async fn transform_request(
        &self,
        physical_region_id: RegionId,
        logical_region_id: RegionId,
        mut request: ScanRequest,
    ) -> Result<ScanRequest> {
        // transform projection
        let physical_projection = match request.projection.as_ref() {
            Some(projection) => {
                self.transform_projection(physical_region_id, logical_region_id, projection)
                    .await?
            }
            None => {
                self.default_projection(physical_region_id, logical_region_id)
                    .await?
            }
        };

        // Rewrite the projection from logical-region schema indices to
        // physical-region schema indices.
        request.projection = Some(physical_projection);

        request
            .filters
            .push(self.table_id_filter(logical_region_id));

        Ok(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(
        &self,
        physical_region_id: RegionId,
        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(|&index| {
                all_logical_columns
                    .get(index)
                    .map(String::as_str)
                    .with_context(|| InvalidRequestSnafu {
                        region_id: logical_region_id,
                        reason: format!("projection index {index} is out of bounds"),
                    })
            })
            .collect::<Result<Vec<_>>>()?;

        // generate physical projection
        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)?;

        for name in projected_logical_names {
            // Safety: logical columns is a strict subset of physical columns
            physical_projection.push(physical_metadata.column_index_by_name(name).unwrap());
        }

        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)
    }

    pub async fn logical_region_metadata(
        &self,
        physical_region_id: RegionId,
        logical_region_id: RegionId,
    ) -> Result<RegionMetadataRef> {
        let logical_columns = self
            .load_logical_columns(physical_region_id, logical_region_id)
            .await?;

        let primary_keys = logical_columns
            .iter()
            .filter_map(|col| {
                if col.semantic_type == SemanticType::Tag {
                    Some(col.column_id)
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        let mut logical_metadata_builder = RegionMetadataBuilder::new(logical_region_id);
        for col in logical_columns {
            logical_metadata_builder.push_column_metadata(col);
        }
        logical_metadata_builder.primary_key(primary_keys);
        let logical_metadata = logical_metadata_builder
            .build()
            .context(InvalidMetadataSnafu)?;

        Ok(Arc::new(logical_metadata))
    }
}

#[cfg(test)]
impl MetricEngineInner {
    pub async fn scan_to_stream(
        &self,
        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)
        }
    }
}

#[cfg(test)]
mod test {
    use std::fmt;

    use api::v1::Rows;
    use common_error::ext::ErrorExt;
    use common_error::status_code::StatusCode;
    use datafusion::physical_plan::DisplayFormatType;
    use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
    use futures_util::TryStreamExt;
    use futures_util::future::try_join_all;
    use mito2::config::MitoConfig;
    use store_api::region_engine::{PrepareRequest, QueryScanContext};
    use store_api::region_request::{RegionFlushRequest, RegionPutRequest, RegionRequest};
    use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector};

    use super::*;
    use crate::config::EngineConfig;
    use crate::test_util::{
        self, TestEnv, alter_logical_region_add_tag_columns, create_logical_region_request,
    };

    struct ScannerDisplay<'a>(&'a dyn store_api::region_engine::RegionScanner);

    impl fmt::Display for ScannerDisplay<'_> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            self.0.fmt_as(DisplayFormatType::Default, f)
        }
    }

    async fn count_scanner_rows(scanner: &mut RegionScannerRef, partitions: usize) -> usize {
        let ranges = scanner
            .properties()
            .partitions
            .iter()
            .flatten()
            .copied()
            .collect::<Vec<_>>();
        let mut prepared = vec![Vec::new(); partitions];
        prepared[0] = ranges;
        scanner
            .prepare(
                PrepareRequest::default()
                    .with_ranges(prepared)
                    .with_target_partitions(partitions),
            )
            .unwrap();

        let metrics = ExecutionPlanMetricsSet::default();
        let context = QueryScanContext::default();
        let streams = (0..partitions)
            .map(|partition| {
                scanner
                    .scan_partition(&context, &metrics, partition)
                    .unwrap()
            })
            .collect::<Vec<_>>();
        try_join_all(streams.into_iter().map(|stream| async move {
            stream
                .try_fold(0, |rows, batch| async move { Ok(rows + batch.num_rows()) })
                .await
        }))
        .await
        .unwrap()
        .into_iter()
        .sum()
    }

    #[tokio::test]
    async fn test_invalid_logical_projection() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();
        let invalid_index = usize::MAX;
        let request = ScanRequest {
            projection: Some(vec![invalid_index]),
            ..Default::default()
        };

        let error =
            match RegionEngine::handle_query(&env.metric(), logical_region_id, request).await {
                Ok(_) => panic!("invalid logical projection unexpectedly succeeded"),
                Err(error) => error,
            };

        assert_eq!(error.status_code(), StatusCode::InvalidArguments);
        assert!(
            error.to_string().contains(&format!(
                "projection index {invalid_index} is out of bounds"
            )),
            "unexpected error: {error}"
        );
    }

    #[tokio::test]
    async fn test_single_logical_region_scan_properties() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let physical_region_id = env.default_physical_region_id();
        let logical_region_id = env.default_logical_region_id();
        let table_id = logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME);
        let table_filter = table_id
            .clone()
            .eq(logical_expr::lit(logical_region_id.table_id()));
        let job_filter = logical_expr::col("job").eq(logical_expr::lit("api"));
        for (filters, logical_region) in [
            (vec![], false),
            (vec![table_filter.clone()], true),
            (
                vec![logical_expr::lit(logical_region_id.table_id()).eq(table_id.clone())],
                true,
            ),
            (vec![job_filter.clone(), table_filter.clone()], true),
            (vec![job_filter.clone().and(table_filter.clone())], true),
            (
                vec![
                    table_filter.clone().and(
                        job_filter
                            .clone()
                            .or(logical_expr::col("job").eq(logical_expr::lit("db"))),
                    ),
                ],
                true,
            ),
            (
                vec![
                    table_filter
                        .clone()
                        .or(table_id.clone().eq(logical_expr::lit(4u32))),
                ],
                false,
            ),
            (vec![table_filter.clone().or(job_filter.clone())], false),
            (
                vec![table_id.clone().not_eq(logical_expr::lit(4u32))],
                false,
            ),
            (vec![table_id.clone().gt(logical_expr::lit(4u32))], false),
            (
                vec![table_id.eq(logical_expr::lit(ScalarValue::UInt32(None)))],
                false,
            ),
            (vec![job_filter], false),
        ] {
            let scanner = env
                .metric()
                .handle_query(
                    physical_region_id,
                    ScanRequest {
                        filters: filters.clone(),
                        distribution: Some(TimeSeriesDistribution::PerSeries),
                        ..Default::default()
                    },
                )
                .await
                .unwrap();
            assert_eq!(
                scanner.properties().is_logical_region(),
                logical_region,
                "{filters:?}"
            );
        }

        let scanner = env
            .metric()
            .handle_query(logical_region_id, ScanRequest::default())
            .await
            .unwrap();
        assert!(scanner.properties().is_logical_region());

        // A table-ID filter alone cannot grant metric guarantees to a Mito scan.
        let scanner = env
            .mito()
            .handle_query(
                utils::to_data_region_id(physical_region_id),
                ScanRequest {
                    filters: vec![table_filter],
                    distribution: Some(TimeSeriesDistribution::PerSeries),
                    ..Default::default()
                },
            )
            .await
            .unwrap();
        assert!(!scanner.properties().is_logical_region());
    }

    #[tokio::test]
    async fn test_transform_scan_req() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;

        let logical_region_id = env.default_logical_region_id();
        let physical_region_id = env.default_physical_region_id();

        // create another logical region
        let logical_region_id2 = RegionId::new(1112345678, 999);
        let create_request =
            create_logical_region_request(&["123", "456", "789"], physical_region_id, "blabla");
        env.metric()
            .handle_request(logical_region_id2, RegionRequest::Create(create_request))
            .await
            .unwrap();

        // add columns to the first logical region
        let alter_request =
            alter_logical_region_add_tag_columns(123456, &["987", "798", "654", "321"]);
        env.metric()
            .handle_request(logical_region_id, RegionRequest::Alter(alter_request))
            .await
            .unwrap();

        // check explicit projection
        let projection = Some(vec![0, 1, 2, 3, 4, 5, 6]);
        let scan_req = ScanRequest {
            projection,
            filters: vec![],
            series_row_selector: Some(TimeSeriesRowSelector::LastRow { after_merge: true }),
            ..Default::default()
        };

        let scan_req = env
            .metric()
            .inner
            .transform_request(physical_region_id, logical_region_id, scan_req)
            .await
            .unwrap();

        assert_eq!(
            scan_req.projection.as_deref().unwrap(),
            &[11, 10, 9, 8, 0, 1, 4]
        );
        assert_eq!(scan_req.filters.len(), 1);
        assert_eq!(
            scan_req.series_row_selector,
            Some(TimeSeriesRowSelector::LastRow { after_merge: true })
        );
        assert_eq!(
            scan_req.filters[0],
            logical_expr::col(DATA_SCHEMA_TABLE_ID_COLUMN_NAME)
                .eq(logical_expr::lit(logical_region_id.table_id()))
        );

        // check default projection
        let scan_req = ScanRequest::default();
        let scan_req = env
            .metric()
            .inner
            .transform_request(physical_region_id, logical_region_id, scan_req)
            .await
            .unwrap();
        assert_eq!(
            scan_req.projection.as_deref().unwrap(),
            &[11, 10, 9, 8, 0, 1, 4]
        );
    }

    #[tokio::test]
    async fn test_two_phase_series_scan_reads_metric_region() {
        let env = TestEnv::with_mito_config(
            "test_two_phase_series_scan",
            MitoConfig {
                experimental_series_scan_v2: true,
                ..Default::default()
            },
            EngineConfig::default(),
        )
        .await;
        env.init_metric_region().await;

        let physical_region_id = env.default_physical_region_id();
        let logical_region_id = env.default_logical_region_id();
        let logical_region_id_2 = RegionId::new(1024, logical_region_id.region_number());
        env.metric()
            .handle_request(
                logical_region_id_2,
                RegionRequest::Create(create_logical_region_request(
                    &["job"],
                    physical_region_id,
                    "test_metric_region_2",
                )),
            )
            .await
            .unwrap();

        let schema = test_util::row_schema_with_tags(&["job"]);
        let put = |rows| {
            RegionRequest::Put(RegionPutRequest {
                skip_wal: false,
                rows: Rows {
                    schema: schema.clone(),
                    rows: test_util::build_rows(1, rows),
                },
                hint: None,
                partition_expr_version: None,
            })
        };
        env.metric()
            .handle_request(logical_region_id, put(3))
            .await
            .unwrap();
        env.metric()
            .handle_request(
                physical_region_id,
                RegionRequest::Flush(RegionFlushRequest::default()),
            )
            .await
            .unwrap();
        env.metric()
            .handle_request(logical_region_id, put(2))
            .await
            .unwrap();
        env.metric()
            .handle_request(logical_region_id_2, put(4))
            .await
            .unwrap();

        let data_region_id = utils::to_data_region_id(physical_region_id);
        let request = ScanRequest {
            distribution: Some(TimeSeriesDistribution::PerSeries),
            ..Default::default()
        };
        let mut scanner = env
            .mito()
            .handle_query(data_region_id, request)
            .await
            .unwrap();

        let rows = tokio::time::timeout(
            std::time::Duration::from_secs(30),
            count_scanner_rows(&mut scanner, 2),
        )
        .await
        .expect("two-phase series scan should not deadlock");
        assert_eq!(7, rows);

        let explain = ScannerDisplay(scanner.as_ref()).to_string();
        assert!(explain.contains("\"mode\":\"two_phase\""), "{explain}");
    }

    #[tokio::test]
    async fn test_series_scan_v2_flag_disables_two_phase_mode() {
        let env = TestEnv::with_mito_config(
            "test_disable_two_phase_series_scan",
            MitoConfig {
                experimental_series_scan_v2: false,
                ..Default::default()
            },
            EngineConfig::default(),
        )
        .await;
        env.init_metric_region().await;

        let data_region_id = utils::to_data_region_id(env.default_physical_region_id());
        let scanner = env
            .mito()
            .handle_query(
                data_region_id,
                ScanRequest {
                    distribution: Some(TimeSeriesDistribution::PerSeries),
                    ..Default::default()
                },
            )
            .await
            .unwrap();

        let explain = ScannerDisplay(scanner.as_ref()).to_string();
        assert!(explain.contains("\"mode\":\"legacy\""), "{explain}");
    }
}