servers/

prom_store.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.

//! prometheus protocol supportings
//! handles prometheus remote_write, remote_read logic
use std::cmp::Ordering;
use std::collections::HashMap;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::Arc;

use api::prom_store::remote::label_matcher::Type as MatcherType;
use api::prom_store::remote::{Label, LabelMatcher, Query, Sample, TimeSeries, WriteRequest};
use api::v1::RowInsertRequests;
use arrow::array::{
    Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, StringArray, StringViewArray,
};
use arrow::compute::kernels::cast as casts;
use arrow::datatypes::{
    DataType, Float64Type, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
    TimestampNanosecondType, UInt32Type,
};
use common_grpc::precision::Precision;
use common_memory_manager::MemoryGuard;
use common_query::native_histogram;
use common_query::prelude::{greptime_timestamp, greptime_value};
use common_recordbatch::{RecordBatch, RecordBatches};
use common_telemetry::{tracing, warn};
use datafusion::dataframe::DataFrame;
use datafusion::prelude::{Expr, ident, lit, regexp_match};
use datafusion_common::ScalarValue;
use datafusion_expr::LogicalPlan;
use snafu::{OptionExt, ResultExt, ensure};
use snap::raw::{Decoder, Encoder};

use crate::error::{self, Result};
use crate::prom_remote_write::REMOTE_WRITE_V1_VERSION;
use crate::request_memory_limiter::ServerMemoryLimiter;
use crate::request_memory_metrics::RequestMemoryMetrics;
use crate::row_writer::{self, MultiTableData};

pub const METRIC_NAME_LABEL: &str = "__name__";
pub const METRIC_NAME_LABEL_BYTES: &[u8] = b"__name__";

/// special label for selecting database name on remote write
pub const DATABASE_LABEL: &str = "x_greptime_database";
pub const DATABASE_LABEL_BYTES: &[u8] = b"x_greptime_database";
pub const DATABASE_LABEL_ALT: &str = "__database__";
pub const DATABASE_LABEL_ALT_BYTES: &[u8] = b"__database__";

/// deprecated, use DATABASE_LABEL instead
#[deprecated(note = "use DATABASE_LABEL instead")]
pub const SCHEMA_LABEL: &str = "__schema__";
#[deprecated(note = "use DATABASE_LABEL_BYTES instead")]
pub const SCHEMA_LABEL_BYTES: &[u8] = b"__schema__";

/// special label for selecting physical table name on remote write
pub const PHYSICAL_TABLE_LABEL: &str = "x_greptime_physical_table";
pub const PHYSICAL_TABLE_LABEL_BYTES: &[u8] = b"x_greptime_physical_table";
pub const PHYSICAL_TABLE_LABEL_ALT: &str = "__physical_table__";
pub const PHYSICAL_TABLE_LABEL_ALT_BYTES: &[u8] = b"__physical_table__";

/// The same as `FIELD_COLUMN_MATCHER` in `promql` crate
pub const FIELD_NAME_LABEL: &str = "__field__";

/// Check if given label is a special label for remote write
#[allow(deprecated)]
pub fn is_remote_write_special_label(label: &str) -> bool {
    label == DATABASE_LABEL
        || label == DATABASE_LABEL_ALT
        || label == PHYSICAL_TABLE_LABEL
        || label == PHYSICAL_TABLE_LABEL_ALT
        || label == SCHEMA_LABEL
}

#[allow(deprecated)]
pub fn is_remote_read_special_label(label: &str) -> bool {
    label == METRIC_NAME_LABEL
        || label == DATABASE_LABEL
        || label == DATABASE_LABEL_ALT
        || label == SCHEMA_LABEL
}

/// Check if given label is a database selection label
#[allow(deprecated)]
pub fn is_database_selection_label(label: &str) -> bool {
    label == DATABASE_LABEL || label == DATABASE_LABEL_ALT || label == SCHEMA_LABEL
}

/// Check if given label is a physical table selection label
pub fn is_physical_table_selection_label(label: &str) -> bool {
    label == PHYSICAL_TABLE_LABEL || label == PHYSICAL_TABLE_LABEL_ALT
}

/// Get table name from remote query
pub fn table_name(q: &Query) -> Result<String> {
    let mut matchers = q
        .matchers
        .iter()
        .filter(|matcher| matcher.name == METRIC_NAME_LABEL);
    let matcher = matchers
        .next()
        .context(error::InvalidPromRemoteRequestSnafu {
            msg: "missing '__name__' label in timeseries",
        })?;

    if matcher.r#type != MatcherType::Eq as i32
        || matcher.value.is_empty()
        || matchers.next().is_some()
    {
        return Err(error::InvalidPromRemoteRequestSnafu {
            msg: "expected exactly one non-empty equality matcher for '__name__'".to_string(),
        }
        .build());
    }

    Ok(matcher.value.clone())
}

/// Extract database selector from a remote read query.
pub fn extract_schema_from_query(query: &Query) -> Option<String> {
    query
        .matchers
        .iter()
        .find(|matcher| {
            is_database_selection_label(&matcher.name) && matcher.r#type == MatcherType::Eq as i32
        })
        .map(|matcher| matcher.value.clone())
}

/// Validate matcher syntax independently of whether the metric exists.
pub fn validate_matchers(matchers: &[LabelMatcher]) -> Result<()> {
    for matcher in matchers {
        let matcher_type = MatcherType::try_from(matcher.r#type).map_err(|e| {
            error::InvalidPromRemoteRequestSnafu {
                msg: format!("invalid LabelMatcher type, decode error: {e}"),
            }
            .build()
        })?;
        if matches!(matcher_type, MatcherType::Re | MatcherType::Nre) {
            // Like Prometheus, validate the raw regex first, then the anchored pattern.
            // We use the same regex engine as query execution, whose syntax differs
            // from Prometheus's: e.g. `(?x) a` is accepted here but rejected by Prometheus.
            regex::Regex::new(&matcher.value)
                .and_then(|_| regex::Regex::new(&anchored_label_regex(&matcher.value)))
                .map_err(|e| {
                    error::InvalidPromRemoteRequestSnafu {
                        msg: format!("invalid regex for label '{}': {e}", matcher.name),
                    }
                    .build()
                })?;
        }
    }
    Ok(())
}

/// Create a DataFrame from a remote Query
#[tracing::instrument(skip_all)]
pub fn query_to_plan(
    dataframe: DataFrame,
    q: &Query,
    timestamp_column_name: &str,
) -> Result<LogicalPlan> {
    let (start_timestamp_ms, end_timestamp_ms) = remote_read_time_range(q);

    let label_matches = &q.matchers;

    let mut conditions = Vec::with_capacity(label_matches.len() + 1);

    conditions
        .push(ident(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms)));
    conditions
        .push(ident(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms)));

    for m in label_matches {
        let name = &m.name;

        if is_remote_read_special_label(name) {
            continue;
        }

        let value = &m.value;
        let m_type = MatcherType::try_from(m.r#type).map_err(|e| {
            error::InvalidPromRemoteRequestSnafu {
                msg: format!("invalid LabelMatcher type, decode error: {e}",),
            }
            .build()
        })?;

        // Prometheus treats an absent label as "".
        let label = if dataframe.schema().has_column_with_unqualified_name(name) {
            ident(name)
        } else {
            lit("")
        };

        // Like PromQL, include NULL when the matcher accepts the empty string.
        // Keep EQ/NEQ on the original column for index recognition.
        let condition = match m_type {
            MatcherType::Eq => {
                if value.is_empty() {
                    label.clone().eq(lit(value)).or(label.is_null())
                } else {
                    label.eq(lit(value))
                }
            }
            MatcherType::Neq => {
                if value.is_empty() {
                    label.not_eq(lit(value))
                } else {
                    label.clone().not_eq(lit(value)).or(label.is_null())
                }
            }
            MatcherType::Re | MatcherType::Nre => {
                let pattern = anchored_label_regex(value);
                let matches_empty = regex::Regex::new(&pattern)
                    .map_err(|e| {
                        error::InvalidPromRemoteRequestSnafu {
                            msg: format!("invalid regex for label '{name}': {e}"),
                        }
                        .build()
                    })?
                    .is_match("");
                let matched = regexp_match(label.clone(), lit(pattern), None);
                if m_type == MatcherType::Re {
                    let condition = matched.is_not_null();
                    if matches_empty {
                        condition.or(label.is_null())
                    } else {
                        condition
                    }
                } else {
                    let condition = matched.is_null();
                    // A NULL input also produces NULL. Keep it only when the
                    // negative matcher accepts the empty string.
                    if matches_empty {
                        condition.and(label.is_not_null())
                    } else {
                        condition
                    }
                }
            }
        };
        conditions.push(condition);
    }

    // Safety: conditions MUST not be empty, reduce always return Some(expr).
    let conditions = conditions.into_iter().reduce(Expr::and).unwrap();

    let dataframe = dataframe
        .filter(conditions)
        .context(error::DataFrameSnafu)?;

    Ok(dataframe.into_parts().1)
}

/// Uses selector time hints without expanding the client's remote-storage range.
/// Nonpositive or reversed hints retain the outer range: proto3 scalar defaults
/// cannot distinguish missing endpoints from explicit zero timestamps. Valid
/// disjoint ranges produce contradictory time predicates and therefore no rows.
fn remote_read_time_range(query: &Query) -> (i64, i64) {
    let outer = (query.start_timestamp_ms, query.end_timestamp_ms);
    let Some(hints) = &query.hints else {
        return outer;
    };
    if hints.start_ms <= 0 || hints.end_ms <= 0 || hints.start_ms > hints.end_ms {
        return outer;
    }

    (outer.0.max(hints.start_ms), outer.1.min(hints.end_ms))
}

/// Prometheus's NewFastRegexMatcher also wraps patterns in ^(?s:...)$:
/// https://github.com/prometheus/prometheus/blob/main/model/labels/regexp.go
fn anchored_label_regex(value: &str) -> String {
    format!("^(?s:{value})$")
}

#[inline]
fn new_label(name: String, value: String) -> Label {
    Label { name, value }
}

fn lit_timestamp_millisecond(ts: i64) -> Expr {
    Expr::Literal(ScalarValue::TimestampMillisecond(Some(ts), None), None)
}

/// Sort timeseries by labels, matching the former `BTreeMap` order.
fn compare_timeseries_labels(left: &[Label], right: &[Label]) -> Ordering {
    let ordering = left.len().cmp(&right.len());
    if ordering != Ordering::Equal {
        return ordering;
    }

    for (left, right) in left.iter().zip(right) {
        let ordering = left.name.cmp(&right.name);
        if ordering != Ordering::Equal {
            return ordering;
        }

        let ordering = left.value.cmp(&right.value);
        if ordering != Ordering::Equal {
            return ordering;
        }
    }

    Ordering::Equal
}

enum LabelValues<'a> {
    Utf8(&'a StringArray),
    LargeUtf8(&'a LargeStringArray),
    Utf8View(&'a StringViewArray),
    DictionaryUtf8 {
        dictionary: &'a DictionaryArray<UInt32Type>,
        values: &'a StringArray,
    },
    Other(Vec<Option<String>>),
}

impl LabelValues<'_> {
    fn value(&self, row: usize) -> Option<&str> {
        match self {
            Self::Utf8(values) => values.is_valid(row).then(|| values.value(row)),
            Self::LargeUtf8(values) => values.is_valid(row).then(|| values.value(row)),
            Self::Utf8View(values) => values.is_valid(row).then(|| values.value(row)),
            Self::DictionaryUtf8 { dictionary, values } => dictionary
                .key(row)
                .and_then(|key| values.is_valid(key).then(|| values.value(key))),
            Self::Other(values) => values.get(row).and_then(Option::as_deref),
        }
    }
}

fn row_labels<'a>(
    columns: &'a [LabelColumn<'a>],
    row: usize,
) -> impl Iterator<Item = (&'a str, &'a str)> {
    columns
        .iter()
        .filter_map(move |column| column.values.value(row).map(|value| (column.name, value)))
}

struct LabelColumn<'a> {
    name: &'a str,
    values: LabelValues<'a>,
}

fn label_columns<'a>(
    recordbatch: &'a RecordBatch,
    timestamp_column_name: &str,
    value_column_name: &str,
) -> Result<Vec<LabelColumn<'a>>> {
    recordbatch
        .schema
        .column_schemas()
        .iter()
        .enumerate()
        .filter(|(_, column_schema)| {
            column_schema.name != timestamp_column_name && column_schema.name != value_column_name
        })
        .map(|(index, column_schema)| {
            let array = recordbatch.column(index);
            let values = match array.data_type() {
                arrow::datatypes::DataType::Utf8 => LabelValues::Utf8(array.as_string::<i32>()),
                arrow::datatypes::DataType::LargeUtf8 => {
                    LabelValues::LargeUtf8(array.as_string::<i64>())
                }
                arrow::datatypes::DataType::Utf8View => {
                    LabelValues::Utf8View(array.as_string_view())
                }
                arrow::datatypes::DataType::Dictionary(key, value)
                    if key.as_ref() == &arrow::datatypes::DataType::UInt32
                        && value.as_ref() == &arrow::datatypes::DataType::Utf8 =>
                {
                    let dictionary = array.as_dictionary::<UInt32Type>();
                    LabelValues::DictionaryUtf8 {
                        dictionary,
                        values: dictionary.values().as_string::<i32>(),
                    }
                }
                _ => {
                    let values = recordbatch.iter_column_as_string(index).collect::<Vec<_>>();
                    ensure!(
                        values.len() == recordbatch.num_rows(),
                        error::InvalidPromRemoteReadQueryResultSnafu {
                            msg: format!(
                                "Cannot convert label column '{}' of datatype {:?} to string",
                                column_schema.name,
                                array.data_type()
                            ),
                        }
                    );
                    LabelValues::Other(values)
                }
            };
            Ok(LabelColumn {
                name: &column_schema.name,
                values,
            })
        })
        .collect()
}

fn hash_timeseries(columns: &[LabelColumn<'_>], row: usize) -> u64 {
    let mut hasher = DefaultHasher::new();

    for (name, value) in row_labels(columns, row) {
        name.hash(&mut hasher);
        value.hash(&mut hasher);
    }

    hasher.finish()
}

fn matches_timeseries(labels: &[Label], columns: &[LabelColumn<'_>], row: usize) -> bool {
    let mut labels = labels.iter().skip(1);
    for (name, value) in row_labels(columns, row) {
        let Some(label) = labels.next() else {
            return false;
        };
        if label.name != name || label.value != value {
            return false;
        }
    }

    labels.next().is_none()
}

fn new_timeseries(table: &str, columns: &[LabelColumn<'_>], row: usize) -> TimeSeries {
    let mut labels = Vec::with_capacity(columns.len() + 1);
    labels.push(new_label(METRIC_NAME_LABEL.to_string(), table.to_string()));

    for (name, value) in row_labels(columns, row) {
        labels.push(new_label(name.to_string(), value.to_string()));
    }

    TimeSeries {
        labels,
        ..Default::default()
    }
}

pub fn recordbatches_to_timeseries(
    table_name: &str,
    timestamp_column_name: &str,
    value_column_name: &str,
    recordbatches: RecordBatches,
) -> Result<Vec<TimeSeries>> {
    let mut timeseries: Vec<TimeSeries> = Vec::new();
    let mut timeseries_by_hash: HashMap<u64, Vec<usize>> = HashMap::new();

    for recordbatch in recordbatches.take() {
        append_recordbatch_to_timeseries(
            table_name,
            timestamp_column_name,
            value_column_name,
            recordbatch,
            &mut timeseries,
            &mut timeseries_by_hash,
        )?;
    }

    for ts in &mut timeseries {
        ts.samples.sort_unstable_by_key(|s| s.timestamp);
        ts.histograms.sort_unstable_by_key(|h| h.timestamp);
    }

    timeseries
        .sort_unstable_by(|left, right| compare_timeseries_labels(&left.labels, &right.labels));
    Ok(timeseries)
}

fn append_recordbatch_to_timeseries(
    table: &str,
    timestamp_column_name: &str,
    value_column_name: &str,
    recordbatch: RecordBatch,
    timeseries: &mut Vec<TimeSeries>,
    timeseries_by_hash: &mut HashMap<u64, Vec<usize>>,
) -> Result<()> {
    let ts_column = recordbatch
        .column_by_name(timestamp_column_name)
        .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
            msg: format!("missing timestamp column '{timestamp_column_name}' in query result"),
        })?;
    // The Prometheus remote read wire format carries millisecond timestamps,
    // while the table's time index can use any time unit (e.g. a metric
    // physical table created with TIMESTAMP(6)). Narrowing floors towards
    // negative infinity, consistent with `Timestamp::convert_to` on the
    // ingestion path; arrow's cast would truncate towards zero, so a
    // pre-epoch sample (e.g. -1001us) would round to -1ms instead of -2ms
    // and disagree with what ingesting the same instant into a millisecond
    // table would have stored.
    let ts_column: ArrayRef = match ts_column.data_type() {
        DataType::Timestamp(TimeUnit::Millisecond, _) => ts_column.clone(),
        // Second -> millisecond is a widening (exact) conversion.
        DataType::Timestamp(TimeUnit::Second, _) => casts::cast(
            ts_column,
            &DataType::Timestamp(TimeUnit::Millisecond, None),
        )
        .map_err(|e| {
            error::InvalidPromRemoteReadQueryResultSnafu {
                msg: format!(
                    "failed to cast timestamp column '{timestamp_column_name}' of datatype {:?} to millisecond: {e}",
                    ts_column.data_type()
                ),
            }
            .build()
        })?,
        DataType::Timestamp(TimeUnit::Microsecond, _) => Arc::new(
            ts_column
                .as_primitive::<TimestampMicrosecondType>()
                .unary::<_, TimestampMillisecondType>(|v| v.div_euclid(1_000)),
        ),
        DataType::Timestamp(TimeUnit::Nanosecond, _) => Arc::new(
            ts_column
                .as_primitive::<TimestampNanosecondType>()
                .unary::<_, TimestampMillisecondType>(|v| v.div_euclid(1_000_000)),
        ),
        _ => {
            return error::InvalidPromRemoteReadQueryResultSnafu {
                msg: format!(
                    "Expect timestamp column of datatype Timestamp(Millisecond), actual {:?}",
                    ts_column.data_type()
                ),
            }
            .fail();
        }
    };
    let ts_column = ts_column
        .as_primitive_opt::<TimestampMillisecondType>()
        .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
            msg: format!(
                "Expect timestamp column of datatype Timestamp(Millisecond), actual {:?}",
                ts_column.data_type()
            ),
        })?;

    let field_column = recordbatch
        .column_by_name(value_column_name)
        .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu {
            msg: format!("missing value column '{value_column_name}' in query result"),
        })?;
    let scalar_column = field_column.as_primitive_opt::<Float64Type>();
    let histogram_column = recordbatch
        .schema
        .column_schema_by_name(value_column_name)
        .filter(|column| native_histogram::is_native_histogram_value_type(&column.data_type))
        .and_then(|_| field_column.as_struct_opt());
    ensure!(
        scalar_column.is_some() || histogram_column.is_some(),
        error::InvalidPromRemoteReadQueryResultSnafu {
            msg: format!(
                "Expect value column of datatype Float64 or canonical native histogram Struct, actual {:?}",
                field_column.data_type()
            ),
        }
    );

    let columns = label_columns(&recordbatch, timestamp_column_name, value_column_name)?;
    let mut previous_timeseries: Option<usize> = None;

    for row in 0..recordbatch.num_rows() {
        let timeseries_index = match previous_timeseries {
            Some(index) if matches_timeseries(&timeseries[index].labels, &columns, row) => index,
            _ => {
                let hash = hash_timeseries(&columns, row);
                let candidates = timeseries_by_hash.entry(hash).or_default();
                match candidates
                    .iter()
                    .copied()
                    .find(|index| matches_timeseries(&timeseries[*index].labels, &columns, row))
                {
                    Some(index) => index,
                    None => {
                        let index = timeseries.len();
                        timeseries.push(new_timeseries(table, &columns, row));
                        candidates.push(index);
                        index
                    }
                }
            }
        };
        previous_timeseries = Some(timeseries_index);

        if ts_column.is_null(row) || field_column.is_null(row) {
            continue;
        }

        let timestamp = ts_column.value(row);
        if let Some(column) = scalar_column {
            timeseries[timeseries_index].samples.push(Sample {
                value: column.value(row),
                timestamp,
            });
        } else if let Some(column) = histogram_column {
            let histogram =
                native_histogram::decode_native_histogram_for_remote_read(column, row, timestamp)
                    .map_err(|error| {
                    error::InvalidPromRemoteReadQueryResultSnafu {
                        msg: error.to_string(),
                    }
                    .build()
                })?;
            if let Some(histogram) = histogram {
                timeseries[timeseries_index].histograms.push(histogram);
            }
        }
    }

    Ok(())
}

pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertRequests, usize)> {
    let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED
        .with_label_values(&["convert", REMOTE_WRITE_V1_VERSION])
        .start_timer();

    let mut multi_table_data = MultiTableData::new();

    for series in &request.timeseries {
        let table_name = &series
            .labels
            .iter()
            .find(|label| {
                // The metric name is a special label
                label.name == METRIC_NAME_LABEL
            })
            .context(error::InvalidPromRemoteRequestSnafu {
                msg: "missing '__name__' label in time-series",
            })?
            .value;

        // The metric name is a special label,
        // num_columns = labels.len() - 1 + 1 (value) + 1 (timestamp)
        let num_columns = series.labels.len() + 1;

        let table_data = multi_table_data.get_or_default_table_data(
            table_name,
            num_columns,
            series.samples.len(),
        );

        // labels
        let kvs = series.labels.iter().filter_map(|label| {
            if label.name == METRIC_NAME_LABEL {
                None
            } else {
                Some((label.name.clone(), label.value.clone()))
            }
        });

        if series.samples.len() == 1 {
            let mut one_row = table_data.alloc_one_row();

            row_writer::write_tags(table_data, kvs, &mut one_row)?;
            // value
            row_writer::write_f64(
                table_data,
                greptime_value(),
                series.samples[0].value,
                &mut one_row,
            )?;
            // timestamp
            row_writer::write_ts_to_millis(
                table_data,
                greptime_timestamp(),
                Some(series.samples[0].timestamp),
                Precision::Millisecond,
                &mut one_row,
            )?;

            table_data.add_row(one_row);
        } else {
            for Sample { value, timestamp } in &series.samples {
                let mut one_row = table_data.alloc_one_row();

                // labels
                let kvs = kvs.clone();
                row_writer::write_tags(table_data, kvs, &mut one_row)?;
                // value
                row_writer::write_f64(table_data, greptime_value(), *value, &mut one_row)?;
                // timestamp
                row_writer::write_ts_to_millis(
                    table_data,
                    greptime_timestamp(),
                    Some(*timestamp),
                    Precision::Millisecond,
                    &mut one_row,
                )?;

                table_data.add_row(one_row);
            }
        }

        if !series.histograms.is_empty() {
            warn!("Native histograms are not supported yet, data ignored");
        }
    }

    Ok(multi_table_data.into_row_insert_requests())
}

/// Hard cap on the decompressed size of a handler-decompressed request body
/// (Prometheus remote read/write, Loki push).
///
/// The HTTP body limit only bounds the *encoded* bytes on the wire. Without
/// this cap a tiny compressed payload (e.g. a five-byte raw-Snappy varint or a
/// few KiB of zstd RLE blocks) could declare/expand to multiple GiB of heap
/// before any protobuf validation runs, so the decompressed size must be
/// bounded on its own.
pub const MAX_DECOMPRESSED_REQUEST_SIZE: usize = 512 * 1024 * 1024;

/// Size of the chunks a zstd body is decompressed and charged in.
const ZSTD_DECOMPRESS_CHUNK: usize = 256 * 1024;

/// Decompressed bytes together with the memory permits charged for them.
/// The permits are released when the buffer is dropped, so the aggregate
/// memory accounting covers the whole lifetime of the decompressed data.
pub(crate) struct ChargedBuffer {
    pub data: Vec<u8>,
    guards: Vec<MemoryGuard<RequestMemoryMetrics>>,
}

impl std::ops::Deref for ChargedBuffer {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

impl std::fmt::Debug for ChargedBuffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ChargedBuffer")
            .field("data_len", &self.data.len())
            .field("guards", &self.guards.len())
            .finish()
    }
}

impl ChargedBuffer {
    pub(crate) fn reserved_bytes(&self) -> u64 {
        self.guards.iter().map(MemoryGuard::granted_bytes).sum()
    }

    /// Creates a buffer with the permits charged for it.
    pub(crate) fn new(data: Vec<u8>, guards: Vec<MemoryGuard<RequestMemoryMetrics>>) -> Self {
        Self { data, guards }
    }

    /// Splits the buffer into the raw bytes and the memory permits charged for
    /// them. Hold the permits as long as the bytes are kept.
    pub(crate) fn into_parts(self) -> (Vec<u8>, Vec<MemoryGuard<RequestMemoryMetrics>>) {
        (self.data, self.guards)
    }
}

/// Decompresses a raw Snappy body after enforcing a hard cap on the decoded
/// size and charging the aggregate request-memory limiter.
///
/// The raw Snappy format prefixes the stream with a varint of the decoded
/// length, so the output size is known (and validated) before the output
/// buffer is allocated by `decompress_vec`.
pub(crate) async fn snappy_decompress_limited(
    buf: &[u8],
    max_decompressed: usize,
    limiter: &ServerMemoryLimiter,
    held_bytes: u64,
) -> Result<ChargedBuffer> {
    let decoded_len =
        snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?;
    ensure!(
        decoded_len <= max_decompressed,
        error::DecompressedBodyTooLargeSnafu {
            size: decoded_len as u64,
            limit: max_decompressed as u64,
        }
    );
    // Reserve the declared size before the allocation happens so concurrent
    // decompressions cannot exceed the aggregate quota.
    limiter.check_request_fits(decoded_len as u64, held_bytes)?;
    let guard = limiter.acquire(decoded_len as u64).await?;
    let mut decoder = Decoder::new();
    let data = decoder
        .decompress_vec(buf)
        .context(error::DecompressSnappyPromRemoteRequestSnafu)?;
    Ok(ChargedBuffer {
        data,
        guards: vec![guard],
    })
}

/// Decompresses a zstd body with a hard cap on the decoded size, charging the
/// aggregate request-memory limiter incrementally while the output is
/// materialized (zstd streams don't declare their decoded size up front).
pub(crate) async fn zstd_decompress_limited(
    buf: &[u8],
    max_decompressed: usize,
    limiter: &ServerMemoryLimiter,
    mut held_bytes: u64,
) -> Result<ChargedBuffer> {
    use std::io::Read;

    let decoder = zstd::stream::read::Decoder::new(buf)
        .context(error::DecompressZstdPromRemoteRequestSnafu)?;
    // Allow reading one byte past the limit so oversized streams are detected
    // instead of being silently truncated.
    let mut limited = decoder.take(max_decompressed as u64 + 1);
    let mut data = Vec::new();
    let mut guards = Vec::new();
    let mut chunk = vec![0u8; ZSTD_DECOMPRESS_CHUNK];
    loop {
        let n = limited
            .read(&mut chunk)
            .context(error::DecompressZstdPromRemoteRequestSnafu)?;
        if n == 0 {
            break;
        }
        ensure!(
            data.len() + n <= max_decompressed,
            error::DecompressedBodyTooLargeSnafu {
                size: (data.len() + n) as u64,
                limit: max_decompressed as u64,
            }
        );
        // Charge each chunk before it is appended so the aggregate quota
        // bounds the memory while it is being materialized.
        limiter.check_request_fits(n as u64, held_bytes)?;
        let guard = limiter.acquire(n as u64).await?;
        held_bytes += guard.granted_bytes();
        guards.push(guard);
        data.extend_from_slice(&chunk[..n]);
    }
    Ok(ChargedBuffer { data, guards })
}

/// Decompresses a raw Snappy body with the hard decoded-size cap enforced,
/// without charging the aggregate request-memory limiter.
///
/// Used for response bodies (remote read replies) and tests; request bodies
/// must use [`snappy_decompress_limited`] instead so their decoded size is
/// also charged to the aggregate quota.
#[inline]
pub fn snappy_decompress(buf: &[u8]) -> Result<Vec<u8>> {
    let decoded_len =
        snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?;
    ensure!(
        decoded_len <= MAX_DECOMPRESSED_REQUEST_SIZE,
        error::DecompressedBodyTooLargeSnafu {
            size: decoded_len as u64,
            limit: MAX_DECOMPRESSED_REQUEST_SIZE as u64,
        }
    );
    let mut decoder = Decoder::new();
    decoder
        .decompress_vec(buf)
        .context(error::DecompressSnappyPromRemoteRequestSnafu)
}

#[inline]
pub fn snappy_compress(buf: &[u8]) -> Result<Vec<u8>> {
    let mut encoder = Encoder::new();
    encoder
        .compress_vec(buf)
        .context(error::CompressPromRemoteRequestSnafu)
}

/// Mock timeseries for test, it is both used in servers and frontend crate
/// So we present it here
pub fn mock_timeseries() -> Vec<TimeSeries> {
    vec![
        TimeSeries {
            labels: vec![
                new_label(METRIC_NAME_LABEL.to_string(), "metric1".to_string()),
                new_label("job".to_string(), "spark".to_string()),
            ],
            samples: vec![
                Sample {
                    value: 1.0f64,
                    timestamp: 1000,
                },
                Sample {
                    value: 2.0f64,
                    timestamp: 2000,
                },
            ],
            ..Default::default()
        },
        TimeSeries {
            labels: vec![
                new_label(METRIC_NAME_LABEL.to_string(), "metric2".to_string()),
                new_label("instance".to_string(), "test_host1".to_string()),
                new_label("idc".to_string(), "z001".to_string()),
            ],
            samples: vec![
                Sample {
                    value: 3.0f64,
                    timestamp: 1000,
                },
                Sample {
                    value: 4.0f64,
                    timestamp: 2000,
                },
            ],
            ..Default::default()
        },
        TimeSeries {
            labels: vec![
                new_label(METRIC_NAME_LABEL.to_string(), "metric3".to_string()),
                new_label("idc".to_string(), "z002".to_string()),
                new_label("app".to_string(), "biz".to_string()),
            ],
            samples: vec![
                Sample {
                    value: 5.0f64,
                    timestamp: 1000,
                },
                Sample {
                    value: 6.0f64,
                    timestamp: 2000,
                },
                Sample {
                    value: 7.0f64,
                    timestamp: 3000,
                },
            ],
            ..Default::default()
        },
    ]
}

/// Add new labels to the mock timeseries.
pub fn mock_timeseries_new_label() -> Vec<TimeSeries> {
    let ts_demo_metrics = TimeSeries {
        labels: vec![
            new_label(METRIC_NAME_LABEL.to_string(), "demo_metrics".to_string()),
            new_label("idc".to_string(), "idc3".to_string()),
            new_label("new_label1".to_string(), "foo".to_string()),
        ],
        samples: vec![Sample {
            value: 42.0,
            timestamp: 3000,
        }],
        ..Default::default()
    };
    let ts_multi_labels = TimeSeries {
        labels: vec![
            new_label(METRIC_NAME_LABEL.to_string(), "metric1".to_string()),
            new_label("idc".to_string(), "idc4".to_string()),
            new_label("env".to_string(), "prod".to_string()),
            new_label("host".to_string(), "host9".to_string()),
            new_label("new_label2".to_string(), "bar".to_string()),
        ],
        samples: vec![Sample {
            value: 99.0,
            timestamp: 4000,
        }],
        ..Default::default()
    };

    vec![ts_demo_metrics, ts_multi_labels]
}

/// Add new labels to the mock timeseries.
pub fn mock_timeseries_special_labels() -> Vec<TimeSeries> {
    let idc3_schema = TimeSeries {
        labels: vec![
            new_label(METRIC_NAME_LABEL.to_string(), "idc3_lo_table".to_string()),
            new_label(DATABASE_LABEL.to_string(), "idc3".to_string()),
            new_label(PHYSICAL_TABLE_LABEL.to_string(), "f1".to_string()),
        ],
        samples: vec![Sample {
            value: 42.0,
            timestamp: 3000,
        }],
        ..Default::default()
    };
    let idc4_schema = TimeSeries {
        labels: vec![
            new_label(
                METRIC_NAME_LABEL.to_string(),
                "idc4_local_table".to_string(),
            ),
            new_label(DATABASE_LABEL.to_string(), "idc4".to_string()),
            new_label(PHYSICAL_TABLE_LABEL.to_string(), "f2".to_string()),
        ],
        samples: vec![Sample {
            value: 99.0,
            timestamp: 4000,
        }],
        ..Default::default()
    };

    vec![idc3_schema, idc4_schema]
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use api::prom_store::remote::LabelMatcher;
    use api::v1::{ColumnDataType, Row, SemanticType};
    use arrow::array::{
        DictionaryArray, Float64Array, StringArray, TimestampMillisecondArray, UInt32Array,
    };
    use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, UInt32Type};
    use common_memory_manager::OnExhaustedPolicy;
    use common_recordbatch::DfRecordBatch;
    use datafusion::prelude::SessionContext;
    use datatypes::data_type::ConcreteDataType;
    use datatypes::schema::{ColumnSchema, Schema};
    use datatypes::vectors::{
        Float64Vector, Int32Vector, StringVector, TimestampMillisecondVector,
    };
    use table::table::adapter::DfTableProviderAdapter;
    use table::test_util::MemTable;

    use super::*;

    const EQ_TYPE: i32 = MatcherType::Eq as i32;
    const NEQ_TYPE: i32 = MatcherType::Neq as i32;
    const RE_TYPE: i32 = MatcherType::Re as i32;

    #[test]
    fn test_validate_matchers_regex_syntax() {
        for op in [MatcherType::Re, MatcherType::Nre] {
            for (value, valid) in [
                ("api|web", true),
                ("(?s)api.web", true),
                ("(?x) api", true),
                ("(?x) api # comment", false),
                ("[", false),
                // Wrapping must not make an invalid original pattern valid.
                ("a)|(?:b", false),
            ] {
                let matcher = LabelMatcher {
                    name: "job".into(),
                    value: value.into(),
                    r#type: op as i32,
                };
                assert_eq!(
                    valid,
                    validate_matchers(&[matcher]).is_ok(),
                    "{op:?}: {value}"
                );
            }
        }
    }

    #[test]
    fn test_table_name() {
        let q = Query {
            start_timestamp_ms: 1000,
            end_timestamp_ms: 2000,
            matchers: vec![],
            ..Default::default()
        };
        let err = table_name(&q).unwrap_err();
        assert!(matches!(err, error::Error::InvalidPromRemoteRequest { .. }));

        let q = Query {
            start_timestamp_ms: 1000,
            end_timestamp_ms: 2000,
            matchers: vec![LabelMatcher {
                name: METRIC_NAME_LABEL.to_string(),
                value: "test".to_string(),
                r#type: EQ_TYPE,
            }],
            ..Default::default()
        };
        assert_eq!("test", table_name(&q).unwrap());

        for matchers in [
            vec![LabelMatcher {
                name: METRIC_NAME_LABEL.to_string(),
                value: "test.*".to_string(),
                r#type: RE_TYPE,
            }],
            vec![LabelMatcher {
                name: METRIC_NAME_LABEL.to_string(),
                value: String::new(),
                r#type: EQ_TYPE,
            }],
            vec![
                LabelMatcher {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "test".to_string(),
                    r#type: EQ_TYPE,
                },
                LabelMatcher {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "other".to_string(),
                    r#type: EQ_TYPE,
                },
            ],
        ] {
            let q = Query {
                matchers,
                ..Default::default()
            };
            assert!(matches!(
                table_name(&q),
                Err(error::Error::InvalidPromRemoteRequest { .. })
            ));
        }
    }

    #[test]
    #[allow(deprecated)]
    fn test_extract_schema_from_query() {
        let query = Query::default();
        assert_eq!(None, extract_schema_from_query(&query));

        for label in [DATABASE_LABEL, DATABASE_LABEL_ALT, SCHEMA_LABEL] {
            let query = Query {
                matchers: vec![LabelMatcher {
                    name: label.to_string(),
                    value: "selected_schema".to_string(),
                    r#type: EQ_TYPE,
                }],
                ..Default::default()
            };
            assert_eq!(
                Some("selected_schema".to_string()),
                extract_schema_from_query(&query)
            );
        }

        let query = Query {
            matchers: vec![LabelMatcher {
                name: DATABASE_LABEL.to_string(),
                value: "selected_schema".to_string(),
                r#type: NEQ_TYPE,
            }],
            ..Default::default()
        };
        assert_eq!(None, extract_schema_from_query(&query));
    }

    #[test]
    fn test_query_to_plan() {
        let q = Query {
            start_timestamp_ms: 1000,
            end_timestamp_ms: 2000,
            matchers: vec![LabelMatcher {
                name: METRIC_NAME_LABEL.to_string(),
                value: "test".to_string(),
                r#type: EQ_TYPE,
            }],
            ..Default::default()
        };

        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                greptime_timestamp(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                true,
            ),
            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
            ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true),
            ColumnSchema::new("job", ConcreteDataType::string_datatype(), true),
            ColumnSchema::new("service.name", ConcreteDataType::string_datatype(), true),
            ColumnSchema::new("Region", ConcreteDataType::string_datatype(), true),
        ]));
        let recordbatch = RecordBatch::new(
            schema,
            vec![
                Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as _,
                Arc::new(Float64Vector::from_vec(vec![3.0])) as _,
                Arc::new(StringVector::from(vec!["host1"])) as _,
                Arc::new(StringVector::from(vec!["job"])) as _,
                Arc::new(StringVector::from(vec!["api"])) as _,
                Arc::new(StringVector::from(vec!["us"])) as _,
            ],
        )
        .unwrap();

        let ctx = SessionContext::new();
        let table = MemTable::table("test", recordbatch);
        let table_provider = Arc::new(DfTableProviderAdapter::new(table));

        let dataframe = ctx.read_table(table_provider.clone()).unwrap();
        let plan = query_to_plan(dataframe, &q, greptime_timestamp()).unwrap();
        let display_string = format!("{}", plan.display_indent());

        let ts_col = greptime_timestamp();
        let expected = format!(
            "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None)\n  TableScan: ?table?",
            ts_col, ts_col
        );
        assert_eq!(expected, display_string);

        let q = Query {
            start_timestamp_ms: 1000,
            end_timestamp_ms: 2000,
            matchers: vec![
                LabelMatcher {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "test".to_string(),
                    r#type: EQ_TYPE,
                },
                LabelMatcher {
                    name: "job".to_string(),
                    value: ".*prom.*".to_string(),
                    r#type: RE_TYPE,
                },
                LabelMatcher {
                    name: "instance".to_string(),
                    value: "localhost".to_string(),
                    r#type: NEQ_TYPE,
                },
                LabelMatcher {
                    name: "service.name".to_string(),
                    value: "api".to_string(),
                    r#type: EQ_TYPE,
                },
                LabelMatcher {
                    name: "Region".to_string(),
                    value: "us".to_string(),
                    r#type: EQ_TYPE,
                },
            ],
            ..Default::default()
        };

        let dataframe = ctx.read_table(table_provider).unwrap();
        let plan = query_to_plan(dataframe, &q, greptime_timestamp()).unwrap();
        let display_string = format!("{}", plan.display_indent());

        let ts_col = greptime_timestamp();
        let expected = format!(
            "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"^(?s:.*prom.*)$\")) IS NOT NULL AND (?table?.instance != Utf8(\"localhost\") OR ?table?.instance IS NULL) AND ?table?.service.name = Utf8(\"api\") AND ?table?.Region = Utf8(\"us\")\n  TableScan: ?table?",
            ts_col, ts_col
        );
        assert_eq!(expected, display_string);
    }

    #[tokio::test]
    async fn test_remote_read_time_hints() {
        use api::prom_store::remote::ReadHints;

        let ctx = SessionContext::new();
        let timestamps = vec![-1, 0, 999, 1000, 1500, 2000, 2001, 3000];
        let batch = DfRecordBatch::try_from_iter(vec![(
            greptime_timestamp(),
            Arc::new(TimestampMillisecondArray::from(timestamps)) as ArrayRef,
        )])
        .unwrap();
        ctx.register_batch("metrics", batch).unwrap();

        for (outer, hinted, expected) in [
            ((1000, 2000), None, vec![1000, 1500, 2000]),
            ((1000, 2000), Some((0, 0)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((0, 1500)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((1500, 0)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((-1, 1500)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((1500, -1)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((2000, 1000)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((1000, 2000)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((1, 3000)), vec![1000, 1500, 2000]),
            ((1000, 2000), Some((1001, 1999)), vec![1500]),
            ((1000, 2000), Some((1500, 3000)), vec![1500, 2000]),
            ((1000, 2000), Some((1, 1500)), vec![1000, 1500]),
            ((1000, 2000), Some((1500, 1500)), vec![1500]),
            ((1000, 2000), Some((2000, 3000)), vec![2000]),
            ((1000, 2000), Some((2001, 3000)), vec![]),
            ((1000, 2000), Some((1, 999)), vec![]),
            ((2000, 1000), Some((1, 3000)), vec![]),
            ((-1, 1000), Some((0, 1000)), vec![-1, 0, 999, 1000]),
        ] {
            let query = Query {
                start_timestamp_ms: outer.0,
                end_timestamp_ms: outer.1,
                hints: hinted.map(|(start_ms, end_ms)| ReadHints {
                    start_ms,
                    end_ms,
                    // These describe evaluation, not sampling or aggregation instructions.
                    step_ms: 10000,
                    range_ms: 300000,
                    func: "rate".into(),
                    grouping: vec!["job".into()],
                    by: true,
                }),
                ..Default::default()
            };
            let plan = query_to_plan(
                ctx.table("metrics").await.unwrap(),
                &query,
                greptime_timestamp(),
            )
            .unwrap();
            let batches = ctx
                .execute_logical_plan(plan)
                .await
                .unwrap()
                .collect()
                .await
                .unwrap();
            let actual: Vec<i64> = batches
                .iter()
                .flat_map(|batch| {
                    batch
                        .column(0)
                        .as_primitive::<TimestampMillisecondType>()
                        .values()
                        .to_vec()
                })
                .collect();
            assert_eq!(expected, actual, "outer={outer:?}, hints={hinted:?}");
        }
    }

    #[tokio::test]
    async fn test_remote_read_label_matcher_semantics() {
        let ctx = SessionContext::new();
        let batch = DfRecordBatch::try_from_iter(vec![
            (
                greptime_timestamp(),
                Arc::new(TimestampMillisecondArray::from(vec![1000; 6])) as ArrayRef,
            ),
            (
                "job",
                Arc::new(StringArray::from(vec![
                    Some("api"),
                    Some("my-api-server"),
                    Some(""),
                    None,
                    Some("web"),
                    Some("api\nweb"),
                ])) as ArrayRef,
            ),
            (
                "id",
                Arc::new(UInt32Array::from(vec![0, 1, 2, 3, 4, 5])) as ArrayRef,
            ),
        ])
        .unwrap();
        ctx.register_batch("metrics", batch).unwrap();
        for (name, op, value, expected) in [
            ("job", MatcherType::Eq, "api", vec![0]),
            ("job", MatcherType::Neq, "api", vec![1, 2, 3, 4, 5]),
            ("job", MatcherType::Eq, "", vec![2, 3]),
            ("job", MatcherType::Neq, "", vec![0, 1, 4, 5]),
            ("job", MatcherType::Re, "api", vec![0]),
            ("job", MatcherType::Nre, "api", vec![1, 2, 3, 4, 5]),
            ("job", MatcherType::Re, "api|web", vec![0, 4]),
            ("job", MatcherType::Nre, "api|web", vec![1, 2, 3, 5]),
            ("job", MatcherType::Re, "", vec![2, 3]),
            ("job", MatcherType::Nre, "", vec![0, 1, 4, 5]),
            ("job", MatcherType::Re, ".*", vec![0, 1, 2, 3, 4, 5]),
            ("job", MatcherType::Nre, ".*", vec![]),
            ("job", MatcherType::Re, ".+", vec![0, 1, 4, 5]),
            ("job", MatcherType::Nre, ".+", vec![2, 3]),
            ("job", MatcherType::Re, "(?m)api", vec![0]),
            ("job", MatcherType::Re, "api.web", vec![5]),
            ("job", MatcherType::Nre, "api.web", vec![0, 1, 2, 3, 4]),
            ("job", MatcherType::Re, "(?-s:api.web)", vec![]),
            ("job", MatcherType::Re, "(api)?", vec![0, 2, 3]),
            ("missing", MatcherType::Eq, "", vec![0, 1, 2, 3, 4, 5]),
            ("missing", MatcherType::Eq, "us", vec![]),
            ("missing", MatcherType::Neq, "us", vec![0, 1, 2, 3, 4, 5]),
            ("missing", MatcherType::Neq, "", vec![]),
            ("missing", MatcherType::Re, ".*", vec![0, 1, 2, 3, 4, 5]),
            ("missing", MatcherType::Re, ".+", vec![]),
            ("missing", MatcherType::Nre, ".*", vec![]),
            ("missing", MatcherType::Nre, ".+", vec![0, 1, 2, 3, 4, 5]),
            ("missing", MatcherType::Re, "", vec![0, 1, 2, 3, 4, 5]),
            ("missing", MatcherType::Nre, "", vec![]),
        ] {
            let query = Query {
                start_timestamp_ms: 1000,
                end_timestamp_ms: 1000,
                matchers: vec![LabelMatcher {
                    name: name.into(),
                    value: value.into(),
                    r#type: op as i32,
                }],
                ..Default::default()
            };
            let plan = query_to_plan(
                ctx.table("metrics").await.unwrap(),
                &query,
                greptime_timestamp(),
            )
            .unwrap_or_else(|err| panic!("{name} {op:?} {value:?}: {err}"));
            let batches = ctx
                .execute_logical_plan(plan)
                .await
                .unwrap()
                .collect()
                .await
                .unwrap();
            let actual: Vec<u32> = batches
                .iter()
                .flat_map(|batch| {
                    batch
                        .column(2)
                        .as_primitive::<UInt32Type>()
                        .values()
                        .to_vec()
                })
                .collect();
            assert_eq!(expected, actual, "{name} {op:?} {value:?}");
        }
    }

    #[tokio::test]
    async fn test_remote_read_matcher_combinations_and_label_types() {
        let labels = StringArray::from(vec![Some("api"), Some("api"), None, Some("api")]);
        for label_type in [
            DataType::Utf8,
            DataType::LargeUtf8,
            DataType::Utf8View,
            DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
        ] {
            let ctx = SessionContext::new();
            let batch = DfRecordBatch::try_from_iter(vec![
                (
                    greptime_timestamp(),
                    Arc::new(TimestampMillisecondArray::from(vec![999, 1000, 2000, 2001]))
                        as ArrayRef,
                ),
                ("Job.Name", casts::cast(&labels, &label_type).unwrap()),
            ])
            .unwrap();
            ctx.register_batch("metrics", batch).unwrap();
            let mut cases = vec![
                (vec![("Job.Name", MatcherType::Re, ".*")], vec![1000, 2000]),
                (vec![("Job.Name", MatcherType::Nre, ".+")], vec![2000]),
            ];
            // Exercise combinations once; other representations only need regex/NULL coverage.
            if label_type == DataType::Utf8 {
                cases.extend([
                    (
                        vec![
                            ("Job.Name", MatcherType::Re, "api|web"),
                            ("Job.Name", MatcherType::Neq, "web"),
                        ],
                        vec![1000],
                    ),
                    (
                        vec![
                            ("Job.Name", MatcherType::Eq, ""),
                            ("Job.Name", MatcherType::Neq, "api"),
                        ],
                        vec![2000],
                    ),
                    (
                        vec![
                            ("Job.Name", MatcherType::Eq, "api"),
                            ("Job.Name", MatcherType::Eq, "web"),
                        ],
                        vec![],
                    ),
                    (
                        vec![
                            ("job.name", MatcherType::Eq, ""),
                            ("Job.Name", MatcherType::Re, ".*"),
                        ],
                        vec![1000, 2000],
                    ),
                ]);
            }
            for (matchers, expected) in cases {
                let mut query = Query {
                    start_timestamp_ms: 1000,
                    end_timestamp_ms: 2000,
                    matchers: matchers
                        .iter()
                        .map(|(name, op, value)| LabelMatcher {
                            name: (*name).into(),
                            value: (*value).into(),
                            r#type: *op as i32,
                        })
                        .collect(),
                    ..Default::default()
                };
                query.matchers.push(LabelMatcher {
                    name: METRIC_NAME_LABEL.into(),
                    value: "metrics".into(),
                    r#type: EQ_TYPE,
                });
                let plan = query_to_plan(
                    ctx.table("metrics").await.unwrap(),
                    &query,
                    greptime_timestamp(),
                )
                .unwrap();
                let batches = ctx
                    .execute_logical_plan(plan)
                    .await
                    .unwrap()
                    .collect()
                    .await
                    .unwrap();
                let actual: Vec<i64> = batches
                    .iter()
                    .flat_map(|batch| {
                        batch
                            .column(0)
                            .as_primitive::<TimestampMillisecondType>()
                            .values()
                            .to_vec()
                    })
                    .collect();
                assert_eq!(expected, actual, "{label_type:?}: {matchers:?}");
            }
        }
    }

    fn column_schemas_with(
        mut kts_iter: Vec<(&str, ColumnDataType, SemanticType)>,
    ) -> Vec<api::v1::ColumnSchema> {
        kts_iter.push((
            greptime_value(),
            ColumnDataType::Float64,
            SemanticType::Field,
        ));
        kts_iter.push((
            greptime_timestamp(),
            ColumnDataType::TimestampMillisecond,
            SemanticType::Timestamp,
        ));

        kts_iter
            .into_iter()
            .map(|(k, t, s)| api::v1::ColumnSchema {
                column_name: k.to_string(),
                datatype: t as i32,
                semantic_type: s as i32,
                ..Default::default()
            })
            .collect()
    }

    fn make_row_with_label(l1: &str, value: f64, timestamp: i64) -> Row {
        Row {
            values: vec![
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::StringValue(l1.to_string())),
                },
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::F64Value(value)),
                },
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
                        timestamp,
                    )),
                },
            ],
        }
    }

    fn make_row_with_2_labels(l1: &str, l2: &str, value: f64, timestamp: i64) -> Row {
        Row {
            values: vec![
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::StringValue(l1.to_string())),
                },
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::StringValue(l2.to_string())),
                },
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::F64Value(value)),
                },
                api::v1::Value {
                    value_data: Some(api::v1::value::ValueData::TimestampMillisecondValue(
                        timestamp,
                    )),
                },
            ],
        }
    }

    #[test]
    fn test_write_request_to_row_insert_exprs() {
        let write_request = WriteRequest {
            timeseries: mock_timeseries(),
            ..Default::default()
        };

        let mut exprs = to_grpc_row_insert_requests(&write_request)
            .unwrap()
            .0
            .inserts;
        exprs.sort_unstable_by(|l, r| l.table_name.cmp(&r.table_name));
        assert_eq!(3, exprs.len());
        assert_eq!("metric1", exprs[0].table_name);
        assert_eq!("metric2", exprs[1].table_name);
        assert_eq!("metric3", exprs[2].table_name);

        let rows = exprs[0].rows.as_ref().unwrap();
        let schema = &rows.schema;
        let rows = &rows.rows;
        assert_eq!(2, rows.len());
        assert_eq!(3, schema.len());
        assert_eq!(
            column_schemas_with(vec![("job", ColumnDataType::String, SemanticType::Tag)]),
            *schema
        );
        assert_eq!(
            &vec![
                make_row_with_label("spark", 1.0, 1000),
                make_row_with_label("spark", 2.0, 2000),
            ],
            rows
        );

        let rows = exprs[1].rows.as_ref().unwrap();
        let schema = &rows.schema;
        let rows = &rows.rows;
        assert_eq!(2, rows.len());
        assert_eq!(4, schema.len());
        assert_eq!(
            column_schemas_with(vec![
                ("instance", ColumnDataType::String, SemanticType::Tag),
                ("idc", ColumnDataType::String, SemanticType::Tag)
            ]),
            *schema
        );
        assert_eq!(
            &vec![
                make_row_with_2_labels("test_host1", "z001", 3.0, 1000),
                make_row_with_2_labels("test_host1", "z001", 4.0, 2000),
            ],
            rows
        );

        let rows = exprs[2].rows.as_ref().unwrap();
        let schema = &rows.schema;
        let rows = &rows.rows;
        assert_eq!(3, rows.len());
        assert_eq!(4, schema.len());
        assert_eq!(
            column_schemas_with(vec![
                ("idc", ColumnDataType::String, SemanticType::Tag),
                ("app", ColumnDataType::String, SemanticType::Tag)
            ]),
            *schema
        );
        assert_eq!(
            &vec![
                make_row_with_2_labels("z002", "biz", 5.0, 1000),
                make_row_with_2_labels("z002", "biz", 6.0, 2000),
                make_row_with_2_labels("z002", "biz", 7.0, 3000),
            ],
            rows
        );
    }

    #[test]
    fn test_native_histogram_remote_read_wire_fixture() {
        use api::greptime_proto::io::prometheus::write::v2::{
            BucketSpan, Histogram,
            histogram::{Count, ZeroCount},
        };
        use api::prom_store::remote::{QueryResult, ReadResponse};
        use common_query::native_histogram::testing::histogram_array;
        use prost::Message;

        let histograms = vec![
            Histogram {
                count: Some(Count::CountInt((1 << 53) + 13)),
                zero_count: Some(ZeroCount::ZeroCountInt((1 << 53) + 1)),
                sum: 10.0,
                schema: 1,
                zero_threshold: 0.001,
                reset_hint: 2,
                start_timestamp: 123,
                positive_spans: vec![
                    BucketSpan {
                        offset: -2,
                        length: 2,
                    },
                    BucketSpan {
                        offset: 3,
                        length: 1,
                    },
                ],
                positive_deltas: vec![1, 2, -1],
                negative_spans: vec![
                    BucketSpan {
                        offset: -3,
                        length: 1,
                    },
                    BucketSpan {
                        offset: 2,
                        length: 2,
                    },
                ],
                negative_deltas: vec![1, 2, -1],
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountFloat(6.0)),
                zero_count: Some(ZeroCount::ZeroCountFloat(0.5)),
                sum: 20.0,
                schema: 2,
                zero_threshold: 0.002,
                reset_hint: 3,
                positive_spans: vec![BucketSpan {
                    offset: 3,
                    length: 2,
                }],
                positive_counts: vec![2.0, 3.5],
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountInt(1)),
                sum: 1.0,
                schema: -53,
                custom_values: vec![1.0, 2.0],
                positive_spans: vec![BucketSpan {
                    offset: 0,
                    length: 1,
                }],
                positive_deltas: vec![1],
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountFloat(3.5)),
                sum: 4.0,
                schema: -53,
                custom_values: vec![1.0, 2.0],
                reset_hint: 3,
                positive_spans: vec![BucketSpan {
                    offset: 0,
                    length: 2,
                }],
                positive_counts: vec![1.0, 2.5],
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountInt(0)),
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountFloat(0.0)),
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountInt(8)),
                sum: f64::from_bits(0x7ff8_0000_0000_0123),
                positive_spans: vec![BucketSpan {
                    offset: 0,
                    length: 2,
                }],
                positive_deltas: vec![3, -1],
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountInt(0)),
                sum: f64::from_bits(0x7ff0_0000_0000_0002),
                ..Default::default()
            },
            Histogram {
                count: Some(Count::CountInt(i64::MAX as u64)),
                zero_count: Some(ZeroCount::ZeroCountInt(i64::MAX as u64)),
                ..Default::default()
            },
        ];
        let values = histograms
            .iter()
            .map(|h| Some(native_histogram::encode_native_histogram(h).unwrap()))
            .collect::<Vec<_>>();
        let array = histogram_array(&values);
        let buckets = array
            .column_by_name("positive_buckets_i64")
            .unwrap()
            .as_list::<i32>()
            .value(0);
        assert_eq!(
            buckets
                .as_primitive::<arrow::datatypes::Int64Type>()
                .values()
                .as_ref(),
            &[1, 3, 2]
        );
        let arrow_schema = Arc::new(ArrowSchema::new(vec![
            Field::new(
                "ts",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("histogram", array.data_type().clone(), true),
        ]));
        let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap());
        let batch = DfRecordBatch::try_new(
            arrow_schema,
            vec![
                Arc::new(TimestampMillisecondArray::from(
                    (1..=9).map(|i| i * 1000).collect::<Vec<_>>(),
                )),
                Arc::new(array),
            ],
        )
        .unwrap();
        let batches = RecordBatches::try_new(
            schema.clone(),
            vec![RecordBatch::from_df_record_batch(schema, batch)],
        )
        .unwrap();
        let response = ReadResponse {
            results: vec![QueryResult {
                timeseries: recordbatches_to_timeseries(
                    "native_histogram_remote_read",
                    "ts",
                    "histogram",
                    batches,
                )
                .unwrap(),
            }],
        };
        let body = snappy_compress(&response.encode_to_vec()).unwrap();
        let mut decoded =
            ReadResponse::decode(snappy_decompress(&body).unwrap().as_slice()).unwrap();
        let mut expected = ReadResponse::decode(
            include_bytes!(
                "../../../tests/compatibility/wire/native_histogram_remote_read/response.pb"
            )
            .as_slice(),
        )
        .unwrap();
        let actual_histograms = &mut decoded.results[0].timeseries[0].histograms;
        assert_eq!(actual_histograms[0].positive_deltas, vec![1, 2, -1]);
        assert_eq!(actual_histograms[0].negative_deltas, vec![1, 2, -1]);
        assert_eq!(actual_histograms[6].positive_deltas, vec![3, -1]);
        for (actual, expected) in actual_histograms
            .iter_mut()
            .zip(&mut expected.results[0].timeseries[0].histograms)
        {
            assert_eq!(actual.sum.to_bits(), expected.sum.to_bits());
            actual.sum = 0.0;
            expected.sum = 0.0;
        }
        assert_eq!(decoded, expected);
        // The optional output is consumed by the pinned Prometheus iterator check.
        if let Some(path) = std::env::var_os("GREPTIME_NATIVE_HISTOGRAM_REMOTE_READ_OUTPUT") {
            std::fs::write(path, body).unwrap();
        }
    }

    #[test]
    fn test_native_histogram_remote_read_groups_and_sorts_batches() {
        use api::greptime_proto::io::prometheus::write::v2::{Histogram, histogram::Count};
        use common_query::native_histogram::testing::histogram_array;

        let integer = native_histogram::encode_native_histogram(&Histogram {
            count: Some(Count::CountInt(0)),
            sum: f64::from_bits(0x7ff0_0000_0000_0002),
            ..Default::default()
        })
        .unwrap();
        let float = native_histogram::encode_native_histogram(&Histogram {
            count: Some(Count::CountFloat(0.0)),
            ..Default::default()
        })
        .unwrap();
        let histograms = histogram_array(&[Some(integer), Some(float.clone()), None, Some(float)]);
        let arrow_schema = Arc::new(ArrowSchema::new(vec![
            Field::new(
                "event_time",
                ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
                true,
            ),
            Field::new("legacy_histogram", histograms.data_type().clone(), true),
            Field::new_dictionary("instance", ArrowDataType::UInt32, ArrowDataType::Utf8, true),
        ]));
        let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap());
        let batch = |timestamps: Vec<Option<i64>>, labels: Vec<Option<u32>>| {
            RecordBatch::from_df_record_batch(
                schema.clone(),
                DfRecordBatch::try_new(
                    arrow_schema.clone(),
                    vec![
                        Arc::new(arrow::array::TimestampMicrosecondArray::from(timestamps)),
                        Arc::new(histograms.clone()),
                        Arc::new(DictionaryArray::<UInt32Type>::new(
                            UInt32Array::from(labels),
                            Arc::new(StringArray::from(vec![Some("b"), Some("a"), None])),
                        )),
                    ],
                )
                .unwrap(),
            )
        };
        let batches = vec![
            batch(
                vec![Some(4000), Some(3000), Some(5000), Some(-1001)],
                vec![Some(0), None, Some(0), Some(0)],
            ),
            batch(
                vec![Some(2000), Some(1000), Some(6000), None],
                vec![Some(0), Some(2), Some(1), Some(0)],
            ),
        ];
        let series = recordbatches_to_timeseries(
            "latency",
            "event_time",
            "legacy_histogram",
            RecordBatches::try_new(schema.clone(), batches).unwrap(),
        )
        .unwrap();
        assert_eq!(series.len(), 3);
        assert_eq!(
            series[0].labels,
            vec![new_label(
                METRIC_NAME_LABEL.to_string(),
                "latency".to_string()
            )]
        );
        assert_eq!(
            series[0]
                .histograms
                .iter()
                .map(|h| h.timestamp)
                .collect::<Vec<_>>(),
            vec![1, 3]
        );
        assert_eq!(series[1].labels[1].value, "a");
        assert!(series[1].histograms.is_empty());
        assert_eq!(series[2].labels[1].value, "b");
        assert_eq!(
            series[2]
                .histograms
                .iter()
                .map(|h| h.timestamp)
                .collect::<Vec<_>>(),
            vec![-2, 2, 4]
        );
        assert!(matches!(
            series[2].histograms[0].count,
            Some(api::prom_store::remote::histogram::Count::CountFloat(0.0))
        ));
        for histogram in &series[2].histograms[1..] {
            assert!(matches!(
                histogram.count,
                Some(api::prom_store::remote::histogram::Count::CountInt(0))
            ));
            assert_eq!(histogram.sum.to_bits(), 0x7ff0_0000_0000_0002);
        }
        assert!(
            series
                .iter()
                .all(|s| s.samples.is_empty()
                    && s.labels.iter().all(|l| l.name != "legacy_histogram"))
        );
    }

    #[test]
    fn test_recordbatches_to_timeseries() {
        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                greptime_timestamp(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                true,
            ),
            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
            ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true),
        ]));

        let recordbatches = RecordBatches::try_new(
            schema.clone(),
            vec![
                RecordBatch::new(
                    schema.clone(),
                    vec![
                        Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as _,
                        Arc::new(Float64Vector::from_vec(vec![3.0])) as _,
                        Arc::new(StringVector::from(vec!["host1"])) as _,
                    ],
                )
                .unwrap(),
                RecordBatch::new(
                    schema,
                    vec![
                        Arc::new(TimestampMillisecondVector::from_vec(vec![2000])) as _,
                        Arc::new(Float64Vector::from_vec(vec![7.0])) as _,
                        Arc::new(StringVector::from(vec!["host2"])) as _,
                    ],
                )
                .unwrap(),
            ],
        )
        .unwrap();

        let timeseries = recordbatches_to_timeseries(
            "metric1",
            greptime_timestamp(),
            greptime_value(),
            recordbatches,
        )
        .unwrap();
        assert_eq!(2, timeseries.len());

        assert_eq!(
            vec![
                Label {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "metric1".to_string(),
                },
                Label {
                    name: "instance".to_string(),
                    value: "host1".to_string(),
                },
            ],
            timeseries[0].labels
        );

        assert_eq!(
            timeseries[0].samples,
            vec![Sample {
                value: 3.0,
                timestamp: 1000,
            }]
        );

        assert_eq!(
            vec![
                Label {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "metric1".to_string(),
                },
                Label {
                    name: "instance".to_string(),
                    value: "host2".to_string(),
                },
            ],
            timeseries[1].labels
        );
        assert_eq!(
            timeseries[1].samples,
            vec![Sample {
                value: 7.0,
                timestamp: 2000,
            }]
        );
    }

    #[test]
    fn test_recordbatches_to_timeseries_merges_across_batches() {
        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                greptime_timestamp(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                true,
            ),
            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
            ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true),
        ]));

        let recordbatches = RecordBatches::try_new(
            schema.clone(),
            vec![
                RecordBatch::new(
                    schema.clone(),
                    vec![
                        Arc::new(TimestampMillisecondVector::from_vec(vec![3000, 1500])) as _,
                        Arc::new(Float64Vector::from_vec(vec![30.0, 15.0])) as _,
                        Arc::new(StringVector::from(vec!["host1", "host2"])) as _,
                    ],
                )
                .unwrap(),
                RecordBatch::new(
                    schema.clone(),
                    vec![
                        Arc::new(TimestampMillisecondVector::from_vec(vec![1000])) as _,
                        Arc::new(Float64Vector::from_vec(vec![10.0])) as _,
                        Arc::new(StringVector::from(vec!["host1"])) as _,
                    ],
                )
                .unwrap(),
                RecordBatch::new(
                    schema,
                    vec![
                        Arc::new(TimestampMillisecondVector::from_vec(vec![2000, 2500])) as _,
                        Arc::new(Float64Vector::from_vec(vec![20.0, 25.0])) as _,
                        Arc::new(StringVector::from(vec!["host1", "host2"])) as _,
                    ],
                )
                .unwrap(),
            ],
        )
        .unwrap();

        let timeseries = recordbatches_to_timeseries(
            "cpu_usage",
            greptime_timestamp(),
            greptime_value(),
            recordbatches,
        )
        .unwrap();

        assert_eq!(2, timeseries.len());

        assert_eq!(
            vec![
                Label {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "cpu_usage".to_string(),
                },
                Label {
                    name: "instance".to_string(),
                    value: "host1".to_string(),
                },
            ],
            timeseries[0].labels
        );
        assert_eq!(
            vec![
                Sample {
                    value: 10.0,
                    timestamp: 1000,
                },
                Sample {
                    value: 20.0,
                    timestamp: 2000,
                },
                Sample {
                    value: 30.0,
                    timestamp: 3000,
                },
            ],
            timeseries[0].samples
        );

        assert_eq!(
            vec![
                Label {
                    name: METRIC_NAME_LABEL.to_string(),
                    value: "cpu_usage".to_string(),
                },
                Label {
                    name: "instance".to_string(),
                    value: "host2".to_string(),
                },
            ],
            timeseries[1].labels
        );
        assert_eq!(
            vec![
                Sample {
                    value: 15.0,
                    timestamp: 1500,
                },
                Sample {
                    value: 25.0,
                    timestamp: 2500,
                },
            ],
            timeseries[1].samples
        );
    }

    #[test]
    fn test_recordbatches_to_timeseries_borrows_and_groups_dictionary_labels() {
        let arrow_schema = Arc::new(ArrowSchema::new(vec![
            Field::new(
                greptime_timestamp(),
                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
                false,
            ),
            Field::new(greptime_value(), ArrowDataType::Float64, false),
            Field::new_dictionary("instance", ArrowDataType::UInt32, ArrowDataType::Utf8, true),
        ]));
        let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap());
        let instance = DictionaryArray::<UInt32Type>::new(
            UInt32Array::from(vec![Some(0), None, Some(1), Some(0), Some(2)]),
            Arc::new(StringArray::from(vec![Some("host2"), Some("host1"), None])),
        );
        let batch = DfRecordBatch::try_new(
            arrow_schema,
            vec![
                Arc::new(TimestampMillisecondArray::from(vec![
                    1000, 2000, 3000, 4000, 5000,
                ])),
                Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0])),
                Arc::new(instance),
            ],
        )
        .unwrap();
        let recordbatch = RecordBatch::from_df_record_batch(schema.clone(), batch);
        let columns = label_columns(&recordbatch, greptime_timestamp(), greptime_value()).unwrap();
        assert!(matches!(
            columns[0].values,
            LabelValues::DictionaryUtf8 { .. }
        ));
        drop(columns);
        let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap();

        let timeseries = recordbatches_to_timeseries(
            "metric1",
            greptime_timestamp(),
            greptime_value(),
            recordbatches,
        )
        .unwrap();

        assert_eq!(3, timeseries.len());
        assert_eq!(
            vec![Label {
                name: METRIC_NAME_LABEL.to_string(),
                value: "metric1".to_string(),
            }],
            timeseries[0].labels
        );
        assert_eq!(
            vec![
                Sample {
                    value: 2.0,
                    timestamp: 2000,
                },
                Sample {
                    value: 5.0,
                    timestamp: 5000,
                },
            ],
            timeseries[0].samples
        );
        assert_eq!("host1", timeseries[1].labels[1].value);
        assert_eq!(
            vec![Sample {
                value: 3.0,
                timestamp: 3000,
            }],
            timeseries[1].samples
        );
        assert_eq!("host2", timeseries[2].labels[1].value);
        assert_eq!(
            vec![
                Sample {
                    value: 1.0,
                    timestamp: 1000,
                },
                Sample {
                    value: 4.0,
                    timestamp: 4000,
                },
            ],
            timeseries[2].samples
        );
    }

    #[test]
    fn test_recordbatches_to_timeseries_groups_non_contiguous_series() {
        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                greptime_timestamp(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                true,
            ),
            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
            ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true),
        ]));
        let recordbatch = RecordBatch::new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _,
                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _,
                Arc::new(StringVector::from(vec!["host2", "host1", "host2"])) as _,
            ],
        )
        .unwrap();

        let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap();
        let timeseries = recordbatches_to_timeseries(
            "metric1",
            greptime_timestamp(),
            greptime_value(),
            recordbatches,
        )
        .unwrap();

        // The result stays sorted by labels as it was with the previous BTreeMap.
        assert_eq!("host1", timeseries[0].labels[1].value);
        assert_eq!("host2", timeseries[1].labels[1].value);
        assert_eq!(
            vec![
                Sample {
                    value: 1.0,
                    timestamp: 1000,
                },
                Sample {
                    value: 3.0,
                    timestamp: 3000,
                },
            ],
            timeseries[1].samples
        );
    }

    #[test]
    fn test_recordbatches_to_timeseries_arrow_label_types_and_nulls() {
        let schema = Arc::new(Schema::new(vec![
            ColumnSchema::new(
                greptime_timestamp(),
                ConcreteDataType::timestamp_millisecond_datatype(),
                true,
            ),
            ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true),
            ColumnSchema::new("instance", ConcreteDataType::large_string_datatype(), true),
            ColumnSchema::new("zone", ConcreteDataType::utf8_view_datatype(), true),
            ColumnSchema::new("shard", ConcreteDataType::int32_datatype(), true),
        ]));
        let recordbatch = RecordBatch::new(
            schema.clone(),
            vec![
                Arc::new(TimestampMillisecondVector::from_vec(vec![1000, 2000, 3000])) as _,
                Arc::new(Float64Vector::from_vec(vec![1.0, 2.0, 3.0])) as _,
                Arc::new(StringVector::from(LargeStringArray::from(vec![
                    "host2", "host1", "host2",
                ]))) as _,
                Arc::new(StringVector::from(StringViewArray::from(vec![
                    Some("west"),
                    None,
                    Some("west"),
                ]))) as _,
                Arc::new(Int32Vector::from_vec(vec![2, 1, 2])) as _,
            ],
        )
        .unwrap();

        let recordbatches = RecordBatches::try_new(schema, vec![recordbatch]).unwrap();
        let timeseries = recordbatches_to_timeseries(
            "metric1",
            greptime_timestamp(),
            greptime_value(),
            recordbatches,
        )
        .unwrap();

        assert_eq!(2, timeseries.len());
        assert_eq!(
            vec![
                new_label(METRIC_NAME_LABEL.to_string(), "metric1".to_string()),
                new_label("instance".to_string(), "host1".to_string()),
                new_label("shard".to_string(), "1".to_string()),
            ],
            timeseries[0].labels
        );
        assert_eq!(
            vec![
                new_label(METRIC_NAME_LABEL.to_string(), "metric1".to_string()),
                new_label("instance".to_string(), "host2".to_string()),
                new_label("zone".to_string(), "west".to_string()),
                new_label("shard".to_string(), "2".to_string()),
            ],
            timeseries[1].labels
        );
        assert_eq!(
            vec![
                Sample {
                    value: 1.0,
                    timestamp: 1000,
                },
                Sample {
                    value: 3.0,
                    timestamp: 3000,
                },
            ],
            timeseries[1].samples
        );
    }

    #[tokio::test]
    async fn test_snappy_varint_bomb_is_rejected_before_allocation() {
        // A five-byte raw Snappy header declaring a ~4 GiB decoded length.
        // Regression test for the compression-bomb finding: the declared size
        // must be rejected by the cap check before any buffer is allocated.
        let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f];
        let limiter = ServerMemoryLimiter::default();

        let err = snappy_decompress_limited(&bomb, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0)
            .await
            .unwrap_err();
        assert!(matches!(
            err,
            crate::error::Error::DecompressedBodyTooLarge { .. }
        ));
        assert_eq!(0, limiter.used_bytes());
    }

    #[tokio::test]
    async fn test_snappy_decompress_limited_roundtrip() {
        let raw = vec![b'a'; 1024 * 1024];
        let mut encoder = snap::raw::Encoder::new();
        let compressed = encoder.compress_vec(&raw).unwrap();
        let limiter = ServerMemoryLimiter::default();

        let decoded =
            snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0)
                .await
                .unwrap();
        assert_eq!(&decoded[..], &raw[..]);
    }

    #[tokio::test]
    async fn test_snappy_charges_and_releases_quota() {
        let raw = vec![b'a'; 64 * 1024];
        let mut encoder = snap::raw::Encoder::new();
        let compressed = encoder.compress_vec(&raw).unwrap();
        let limiter = ServerMemoryLimiter::new(128 * 1024, OnExhaustedPolicy::Fail);

        let decoded =
            snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0)
                .await
                .unwrap();
        // Rounding to the 1 KiB permit granularity.
        assert!(limiter.used_bytes() >= 64 * 1024);
        assert!(limiter.used_bytes() <= 65 * 1024);
        drop(decoded);
        assert_eq!(0, limiter.used_bytes());
    }

    #[tokio::test]
    async fn test_snappy_rejects_when_quota_exhausted() {
        let raw = vec![b'a'; 64 * 1024];
        let mut encoder = snap::raw::Encoder::new();
        let compressed = encoder.compress_vec(&raw).unwrap();
        // Quota smaller than the decoded size.
        let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail);

        let err =
            snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0)
                .await
                .unwrap_err();
        assert!(matches!(
            err,
            crate::error::Error::MemoryLimitExceeded { .. }
        ));
        assert_eq!(0, limiter.used_bytes(), "failed charge must be released");
    }

    #[tokio::test]
    async fn test_zstd_at_and_over_the_cap() {
        let limiter = ServerMemoryLimiter::default();
        let small_cap = 8 * 1024;

        // Exactly at the cap: accepted.
        let raw = vec![0u8; small_cap];
        let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap();
        let decoded = zstd_decompress_limited(&compressed, small_cap, &limiter, 0)
            .await
            .unwrap();
        assert_eq!(decoded.len(), small_cap);

        // One byte past the cap: rejected.
        let raw = vec![0u8; small_cap + 1];
        let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap();
        let err = zstd_decompress_limited(&compressed, small_cap, &limiter, 0)
            .await
            .unwrap_err();
        assert!(matches!(
            err,
            crate::error::Error::DecompressedBodyTooLarge { .. }
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn test_zstd_charges_incrementally() {
        // 64 KiB wire that expands to 2 MiB must be charged chunk by chunk
        // against a quota that only fits the expanded size.
        let raw = vec![0u8; 2 * 1024 * 1024];
        let compressed = zstd::stream::encode_all(&raw[..], 19).unwrap();
        assert!(compressed.len() < 16 * 1024);
        let limiter = ServerMemoryLimiter::new(4 * 1024 * 1024, OnExhaustedPolicy::Fail);

        let decoded =
            zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0)
                .await
                .unwrap();
        assert_eq!(decoded.len(), raw.len());
        assert!(limiter.used_bytes() >= raw.len() as u64);
        drop(decoded);
        assert_eq!(0, limiter.used_bytes());
        // The next chunk cannot fit alongside this request's earlier chunks.
        let limiter = ServerMemoryLimiter::new(
            (2 * ZSTD_DECOMPRESS_CHUNK) as u64,
            OnExhaustedPolicy::default(),
        );
        let started = tokio::time::Instant::now();
        let error =
            zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter, 0)
                .await
                .unwrap_err();
        assert!(matches!(error, error::Error::MemoryLimitExceeded { .. }));
        assert_eq!(started.elapsed(), std::time::Duration::ZERO);
        assert_eq!(limiter.used_bytes(), 0);
    }

    #[test]
    fn test_snappy_response_decompress_is_capped() {
        // The response-side (uncharged) helper must enforce the cap too.
        let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f];
        let err = snappy_decompress(&bomb).unwrap_err();
        assert!(matches!(
            err,
            crate::error::Error::DecompressedBodyTooLarge { .. }
        ));
    }
}