common_recordbatch/

adapter.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::fmt::{self, Display};
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::task::{Context, Poll};

use common_base::readable_size::ReadableSize;
use common_telemetry::tracing::{Span, info_span};
use common_time::util::format_nanoseconds_human_readable;
use datafusion::arrow::compute::cast;
use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::Expr;
use datafusion::logical_expr::utils::conjunction;
use datafusion::physical_expr::create_physical_expr;
use datafusion::physical_plan::metrics::{BaselineMetrics, MetricValue};
use datafusion::physical_plan::{
    DisplayFormatType, ExecutionPlan, ExecutionPlanVisitor, PhysicalExpr,
    RecordBatchStream as DfRecordBatchStream, accept,
};
use datafusion_common::arrow::error::ArrowError;
use datafusion_common::{DataFusionError, ToDFSchema};
use datatypes::arrow::array::Array;
use datatypes::arrow::datatypes::DataType as ArrowDataType;
use datatypes::schema::{ColumnExtType, Schema, SchemaRef};
use futures::ready;
use jsonb;
use pin_project::pin_project;
use snafu::ResultExt;

use crate::error::{self, Result};
use crate::filter::batch_filter;
use crate::{
    DfRecordBatch, DfSendableRecordBatchStream, OrderOption, RecordBatch, RecordBatchStream,
    SendableRecordBatchStream, Stream,
};

const REGION_SCAN_EXEC_NAME: &str = "RegionScanExec";

type FutureStream =
    Pin<Box<dyn std::future::Future<Output = Result<SendableRecordBatchStream>> + Send>>;

/// Casts the `RecordBatch`es of `stream` against the `output_schema`.
#[pin_project]
pub struct RecordBatchStreamTypeAdapter<T, E> {
    #[pin]
    stream: T,
    projected_schema: DfSchemaRef,
    projection: Vec<usize>,
    predicate: Option<Arc<dyn PhysicalExpr>>,
    phantom: PhantomData<E>,
}

impl<T, E> RecordBatchStreamTypeAdapter<T, E>
where
    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
    E: std::error::Error + Send + Sync + 'static,
{
    pub fn new(projected_schema: DfSchemaRef, stream: T, projection: Option<Vec<usize>>) -> Self {
        let projection = if let Some(projection) = projection {
            projection
        } else {
            (0..projected_schema.fields().len()).collect()
        };

        Self {
            stream,
            projected_schema,
            projection,
            predicate: None,
            phantom: Default::default(),
        }
    }

    pub fn with_filter(mut self, filters: Vec<Expr>) -> Result<Self> {
        let filters = if let Some(expr) = conjunction(filters) {
            let df_schema = self
                .projected_schema
                .clone()
                .to_dfschema_ref()
                .context(error::PhysicalExprSnafu)?;

            let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
                .context(error::PhysicalExprSnafu)?;
            Some(filters)
        } else {
            None
        };
        self.predicate = filters;
        Ok(self)
    }
}

impl<T, E> DfRecordBatchStream for RecordBatchStreamTypeAdapter<T, E>
where
    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
    E: std::error::Error + Send + Sync + 'static,
{
    fn schema(&self) -> DfSchemaRef {
        self.projected_schema.clone()
    }
}

impl<T, E> Stream for RecordBatchStreamTypeAdapter<T, E>
where
    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
    E: std::error::Error + Send + Sync + 'static,
{
    type Item = DfResult<DfRecordBatch>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();

        let batch = futures::ready!(this.stream.poll_next(cx))
            .map(|r| r.map_err(|e| DataFusionError::External(Box::new(e))));

        let projected_schema = this.projected_schema.clone();
        let projection = this.projection.clone();
        let predicate = this.predicate.clone();

        let batch = batch.map(|b| {
            b.and_then(|b| {
                let projected_column = b.project(&projection)?;
                if projected_column.schema().fields.len() != projected_schema.fields.len() {
                   return Err(DataFusionError::ArrowError(Box::new(ArrowError::SchemaError(format!(
                        "Trying to cast a RecordBatch into an incompatible schema. RecordBatch: {}, Target: {}",
                        projected_column.schema(),
                        projected_schema,
                    ))), None));
                }

                let mut columns = Vec::with_capacity(projected_schema.fields.len());
                for (idx,field) in projected_schema.fields.iter().enumerate() {
                    let column = projected_column.column(idx);
                    let extype = field.metadata().get("greptime:type").and_then(|s| ColumnExtType::from_str(s).ok());
                    let output = custom_cast(&column, field.data_type(), extype)?;
                    columns.push(output)
                }
                let record_batch = DfRecordBatch::try_new(projected_schema, columns)?;
                let record_batch = if let Some(predicate) = predicate {
                    batch_filter(&record_batch, &predicate)?
                } else {
                    record_batch
                };
                Ok(record_batch)
            })
        });

        Poll::Ready(batch)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.stream.size_hint()
    }
}

/// Greptime SendableRecordBatchStream -> DataFusion RecordBatchStream.
/// The reverse one is [RecordBatchStreamAdapter].
pub struct DfRecordBatchStreamAdapter {
    stream: SendableRecordBatchStream,
}

impl DfRecordBatchStreamAdapter {
    pub fn new(stream: SendableRecordBatchStream) -> Self {
        Self { stream }
    }
}

impl DfRecordBatchStream for DfRecordBatchStreamAdapter {
    fn schema(&self) -> DfSchemaRef {
        self.stream.schema().arrow_schema().clone()
    }
}

impl Stream for DfRecordBatchStreamAdapter {
    type Item = DfResult<DfRecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.stream).poll_next(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(Some(recordbatch)) => match recordbatch {
                Ok(recordbatch) => Poll::Ready(Some(Ok(recordbatch.into_df_record_batch()))),
                Err(e) => Poll::Ready(Some(Err(DataFusionError::External(Box::new(e))))),
            },
            Poll::Ready(None) => Poll::Ready(None),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.stream.size_hint()
    }
}

/// DataFusion [SendableRecordBatchStream](DfSendableRecordBatchStream) -> Greptime [RecordBatchStream].
/// The reverse one is [DfRecordBatchStreamAdapter].
/// It can collect metrics from DataFusion execution plan.
pub struct RecordBatchStreamAdapter {
    schema: SchemaRef,
    stream: DfSendableRecordBatchStream,
    metrics: Option<BaselineMetrics>,
    /// Aggregated plan-level metrics. Resolved after an [ExecutionPlan] is finished.
    metrics_2: Metrics,
    query_load_region_id: Option<u64>,
    query_stat_counters: Option<RegionQueryStatCounters>,
    /// Display plan and metrics in verbose mode.
    explain_verbose: bool,
    span: Span,
}

/// Query statistic counters owned by a region.
#[derive(Debug, Clone)]
pub struct RegionQueryStatCounters {
    /// The total query CPU time in nanoseconds.
    pub query_cpu_time: Arc<AtomicU64>,
    /// The total scanned bytes.
    pub query_scanned_bytes: Arc<AtomicU64>,
}

/// Json encoded metrics. Contains metric from a whole plan tree.
enum Metrics {
    Unavailable,
    Unresolved(Arc<dyn ExecutionPlan>),
    PartialResolved(Arc<dyn ExecutionPlan>, RecordBatchMetrics),
    Resolved(RecordBatchMetrics),
}

impl RecordBatchStreamAdapter {
    pub fn try_new(stream: DfSendableRecordBatchStream) -> Result<Self> {
        let schema =
            Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
        Ok(Self {
            schema,
            stream,
            metrics: None,
            metrics_2: Metrics::Unavailable,
            query_load_region_id: None,
            query_stat_counters: None,
            explain_verbose: false,
            span: Span::current(),
        })
    }

    pub fn try_new_with_span(stream: DfSendableRecordBatchStream, span: Span) -> Result<Self> {
        let schema =
            Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
        let subspan = info_span!(parent: &span, "RecordBatchStreamAdapter");
        Ok(Self {
            schema,
            stream,
            metrics: None,
            metrics_2: Metrics::Unavailable,
            query_load_region_id: None,
            query_stat_counters: None,
            explain_verbose: false,
            span: subspan,
        })
    }

    pub fn set_metrics2(&mut self, plan: Arc<dyn ExecutionPlan>) {
        self.metrics_2 = Metrics::Unresolved(plan)
    }

    fn record_query_stats_on_drop(&self) {
        let Some(counters) = &self.query_stat_counters else {
            return;
        };

        match &self.metrics_2 {
            Metrics::Unresolved(df_plan) => {
                let metrics = collect_lightweight_query_load_metrics(
                    df_plan.as_ref(),
                    self.query_load_region_id,
                );
                record_query_stats(counters, &metrics);
            }
            Metrics::PartialResolved(_, metrics) | Metrics::Resolved(metrics) => {
                record_query_stats(counters, metrics);
            }
            Metrics::Unavailable => {}
        }
    }

    pub fn set_query_load_region_id(&mut self, region_id: Option<u64>) {
        self.query_load_region_id = region_id;
    }

    pub fn set_query_stat_counters(&mut self, counters: Option<RegionQueryStatCounters>) {
        self.query_stat_counters = counters;
    }

    /// Set the verbose mode for displaying plan and metrics.
    pub fn set_explain_verbose(&mut self, verbose: bool) {
        self.explain_verbose = verbose;
    }

    fn collect_plan_metrics(&self, df_plan: &Arc<dyn ExecutionPlan>) -> RecordBatchMetrics {
        collect_full_metrics(
            df_plan.as_ref(),
            self.explain_verbose,
            self.query_load_region_id,
        )
    }

    fn collect_partial_metrics(
        df_plan: &dyn ExecutionPlan,
        explain_verbose: bool,
        query_load_region_id: Option<u64>,
    ) -> RecordBatchMetrics {
        if explain_verbose {
            collect_full_metrics(df_plan, explain_verbose, query_load_region_id)
        } else {
            collect_lightweight_query_load_metrics(df_plan, query_load_region_id)
        }
    }

    fn update_plan_metrics(&mut self, final_metrics: bool) {
        if final_metrics {
            let df_plan = match &self.metrics_2 {
                Metrics::Unresolved(df_plan) | Metrics::PartialResolved(df_plan, _) => {
                    df_plan.clone()
                }
                Metrics::Unavailable | Metrics::Resolved(_) => return,
            };
            let metrics = self.collect_plan_metrics(&df_plan);
            self.metrics_2 = Metrics::Resolved(metrics);
        } else {
            let explain_verbose = self.explain_verbose;
            let query_load_region_id = self.query_load_region_id;
            match &mut self.metrics_2 {
                Metrics::Unresolved(df_plan) => {
                    let df_plan = df_plan.clone();
                    let metrics = Self::collect_partial_metrics(
                        df_plan.as_ref(),
                        explain_verbose,
                        query_load_region_id,
                    );
                    self.metrics_2 = Metrics::PartialResolved(df_plan, metrics);
                }
                Metrics::PartialResolved(df_plan, metrics) => {
                    *metrics = Self::collect_partial_metrics(
                        df_plan.as_ref(),
                        explain_verbose,
                        query_load_region_id,
                    );
                }
                Metrics::Unavailable | Metrics::Resolved(_) => {}
            }
        }
    }
}

/// Extracts total `output_bytes` from region scan plan nodes.
pub fn region_scan_output_bytes(metrics: &RecordBatchMetrics) -> usize {
    metrics
        .plan_metrics
        .iter()
        .filter(|pm| pm.plan_name == REGION_SCAN_EXEC_NAME)
        .flat_map(|pm| &pm.metrics)
        .filter_map(|(name, value)| (name == "output_bytes").then_some(*value))
        .sum()
}

fn record_query_stats(counters: &RegionQueryStatCounters, metrics: &RecordBatchMetrics) {
    counters
        .query_cpu_time
        .fetch_add(metrics.elapsed_compute as u64, Ordering::Relaxed);
    counters
        .query_scanned_bytes
        .fetch_add(region_scan_output_bytes(metrics) as u64, Ordering::Relaxed);
}

/// Collects the complete plan metrics used by terminal metrics and verbose analyze output.
fn collect_full_metrics(
    df_plan: &dyn ExecutionPlan,
    explain_verbose: bool,
    query_load_region_id: Option<u64>,
) -> RecordBatchMetrics {
    let mut metric_collector = MetricCollector::new(explain_verbose);
    accept(df_plan, &mut metric_collector).unwrap();
    metric_collector.record_batch_metrics.query_load_region_id = query_load_region_id;
    metric_collector.record_batch_metrics
}

/// Collects the minimal metrics needed for query-load reporting before EOF.
///
/// This intentionally avoids [`MetricCollector`]'s full per-plan aggregation,
/// sorting, and plan formatting work on the normal query hot path. The result
/// is still enough for early-stop/cancellation paths to report per-region CPU
/// time, scanned bytes, and physical-region attribution.
fn collect_lightweight_query_load_metrics(
    df_plan: &dyn ExecutionPlan,
    query_load_region_id: Option<u64>,
) -> RecordBatchMetrics {
    let mut metrics = RecordBatchMetrics {
        query_load_region_id,
        ..Default::default()
    };
    collect_lightweight_query_load_metrics_inner(df_plan, 0, &mut metrics);
    metrics
}

/// Recursively walks the physical plan and reads raw metric values without
/// formatting plan nodes.
fn collect_lightweight_query_load_metrics_inner(
    df_plan: &dyn ExecutionPlan,
    level: usize,
    record_batch_metrics: &mut RecordBatchMetrics,
) {
    let is_region_scan = df_plan.name() == REGION_SCAN_EXEC_NAME;
    let mut region_scan_output_bytes = None;

    if let Some(metrics) = df_plan.metrics() {
        for metric in metrics.iter() {
            let value = metric.value();
            match value {
                MetricValue::ElapsedCompute(elapsed_compute) => {
                    record_batch_metrics.elapsed_compute += elapsed_compute.value();
                }
                MetricValue::CurrentMemoryUsage(memory_usage) => {
                    record_batch_metrics.memory_usage += memory_usage.value();
                }
                _ => {}
            }

            if is_region_scan && value.name() == "output_bytes" {
                *region_scan_output_bytes.get_or_insert(0) += value.as_usize();
            }
        }
    }

    if let Some(output_bytes) = region_scan_output_bytes {
        record_batch_metrics.plan_metrics.push(PlanMetrics {
            plan: df_plan.name().to_string(),
            plan_name: df_plan.name().to_string(),
            level,
            metrics: vec![("output_bytes".to_string(), output_bytes)],
        });
    }

    for child in df_plan.children() {
        collect_lightweight_query_load_metrics_inner(
            child.as_ref(),
            level + 1,
            record_batch_metrics,
        );
    }
}

impl RecordBatchStream for RecordBatchStreamAdapter {
    fn name(&self) -> &str {
        "RecordBatchStreamAdapter"
    }

    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn metrics(&self) -> Option<RecordBatchMetrics> {
        match &self.metrics_2 {
            Metrics::Unresolved(df_plan) => {
                if self.explain_verbose {
                    Some(self.collect_plan_metrics(df_plan))
                } else {
                    None
                }
            }
            Metrics::PartialResolved(df_plan, metrics) => Some(if self.explain_verbose {
                self.collect_plan_metrics(df_plan)
            } else {
                metrics.clone()
            }),
            Metrics::Resolved(metrics) => Some(metrics.clone()),
            Metrics::Unavailable => None,
        }
    }

    fn output_ordering(&self) -> Option<&[OrderOption]> {
        None
    }
}

impl Stream for RecordBatchStreamAdapter {
    type Item = Result<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let timer = self
            .metrics
            .as_ref()
            .map(|m| m.elapsed_compute().clone())
            .unwrap_or_default();
        let _guard = timer.timer();
        let poll_span = info_span!(parent: &self.span, "poll_next");
        let _entered = poll_span.enter();
        match Pin::new(&mut self.stream).poll_next(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(Some(df_record_batch)) => {
                let df_record_batch = df_record_batch?;
                self.update_plan_metrics(false);
                Poll::Ready(Some(Ok(RecordBatch::from_df_record_batch(
                    self.schema(),
                    df_record_batch,
                ))))
            }
            Poll::Ready(None) => {
                self.update_plan_metrics(true);
                Poll::Ready(None)
            }
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.stream.size_hint()
    }
}

impl Drop for RecordBatchStreamAdapter {
    fn drop(&mut self) {
        self.record_query_stats_on_drop();
    }
}

/// An [ExecutionPlanVisitor] to collect metrics from a [ExecutionPlan].
pub struct MetricCollector {
    current_level: usize,
    pub record_batch_metrics: RecordBatchMetrics,
    verbose: bool,
}

impl MetricCollector {
    pub fn new(verbose: bool) -> Self {
        Self {
            current_level: 0,
            record_batch_metrics: RecordBatchMetrics::default(),
            verbose,
        }
    }
}

impl ExecutionPlanVisitor for MetricCollector {
    type Error = !;

    fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
        // skip if no metric available
        let Some(metric) = plan.metrics() else {
            self.record_batch_metrics.plan_metrics.push(PlanMetrics {
                plan: plan.name().to_string(),
                plan_name: plan.name().to_string(),
                level: self.current_level,
                metrics: vec![],
            });
            self.current_level += 1;
            return Ok(true);
        };

        // scrape plan metrics
        let metric = metric
            .aggregate_by_name()
            .sorted_for_display()
            .timestamps_removed();
        let mut plan_metric = PlanMetrics {
            plan: one_line(plan, self.verbose).to_string(),
            plan_name: plan.name().to_string(),
            level: self.current_level,
            metrics: Vec::with_capacity(metric.iter().size_hint().0),
        };
        for m in metric.iter() {
            plan_metric
                .metrics
                .push((m.value().name().to_string(), m.value().as_usize()));

            // aggregate high-level metrics
            match m.value() {
                MetricValue::ElapsedCompute(ec) => {
                    self.record_batch_metrics.elapsed_compute += ec.value()
                }
                MetricValue::CurrentMemoryUsage(m) => {
                    self.record_batch_metrics.memory_usage += m.value()
                }
                _ => {}
            }
        }
        self.record_batch_metrics.plan_metrics.push(plan_metric);

        self.current_level += 1;
        Ok(true)
    }

    fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
        self.current_level -= 1;
        Ok(true)
    }
}

/// Returns a single-line summary of the root of the plan.
/// If the `verbose` flag is set, it will display detailed information about the plan.
fn one_line(plan: &dyn ExecutionPlan, verbose: bool) -> impl fmt::Display + '_ {
    struct Wrapper<'a> {
        plan: &'a dyn ExecutionPlan,
        format_type: DisplayFormatType,
    }

    impl fmt::Display for Wrapper<'_> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            self.plan.fmt_as(self.format_type, f)?;
            writeln!(f)
        }
    }

    let format_type = if verbose {
        DisplayFormatType::Verbose
    } else {
        DisplayFormatType::Default
    };
    Wrapper { plan, format_type }
}

/// [`RecordBatchMetrics`] carrys metrics value
/// from datanode to frontend through gRPC
#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct RecordBatchMetrics {
    // High-level aggregated metrics
    /// CPU consumption in nanoseconds
    pub elapsed_compute: usize,
    /// Memory used by the plan in bytes
    pub memory_usage: usize,
    // Detailed per-plan metrics
    /// An ordered list of plan metrics, from top to bottom in post-order.
    pub plan_metrics: Vec<PlanMetrics>,
    /// Region id that should receive query-load metrics for this scan.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub query_load_region_id: Option<u64>,
    /// Per-region watermark for incremental-read checkpoint advancement.
    ///
    /// The watermark is the latest sequence (`seq`) this query round safely read
    /// for each participating region. Flow uses it to decide where the next
    /// incremental round can resume.
    ///
    /// - `Some(seq)`: the query proved it safely read up to `seq`; downstream
    ///   may advance the checkpoint to this value.
    /// - `None`: the region participated but the query could not prove a safe
    ///   read upper-bound, so the checkpoint must not advance for this region.
    ///
    /// Omitted when empty for backward compatibility.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub region_watermarks: Vec<RegionWatermarkEntry>,
}

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct RegionWatermarkEntry {
    pub region_id: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub watermark: Option<u64>,
}

/// Determines if a metric name represents a time measurement that should be formatted.
fn is_time_metric(metric_name: &str) -> bool {
    metric_name.contains("elapsed") || metric_name.contains("time") || metric_name.contains("cost")
}

/// Determines if a metric name represents a bytes measurement that should be formatted.
fn is_bytes_metric(metric_name: &str) -> bool {
    metric_name.contains("bytes") || metric_name.contains("mem")
}

fn format_bytes_human_readable(bytes: usize) -> String {
    format!("{}", ReadableSize(bytes as u64))
}

/// Only display `plan_metrics` with indent `  ` (2 spaces).
impl Display for RecordBatchMetrics {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for metric in &self.plan_metrics {
            write!(
                f,
                "{:indent$}{} metrics=[",
                " ",
                metric.plan.trim_end(),
                indent = metric.level * 2,
            )?;
            for (label, value) in &metric.metrics {
                if is_time_metric(label) {
                    write!(
                        f,
                        "{}: {}, ",
                        label,
                        format_nanoseconds_human_readable(*value),
                    )?;
                } else if is_bytes_metric(label) {
                    write!(f, "{}: {}, ", label, format_bytes_human_readable(*value),)?;
                } else {
                    write!(f, "{}: {}, ", label, value)?;
                }
            }
            writeln!(f, "]")?;
        }

        Ok(())
    }
}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct PlanMetrics {
    /// The plan name
    pub plan: String,
    /// The stable execution plan name.
    #[serde(default)]
    pub plan_name: String,
    /// The level of the plan, starts from 0
    pub level: usize,
    /// An ordered key-value list of metrics.
    /// Key is metric label and value is metric value.
    pub metrics: Vec<(String, usize)>,
}

enum AsyncRecordBatchStreamAdapterState {
    Uninit(FutureStream),
    Ready(SendableRecordBatchStream),
    Failed,
}

pub struct AsyncRecordBatchStreamAdapter {
    schema: SchemaRef,
    state: AsyncRecordBatchStreamAdapterState,
}

impl AsyncRecordBatchStreamAdapter {
    pub fn new(schema: SchemaRef, stream: FutureStream) -> Self {
        Self {
            schema,
            state: AsyncRecordBatchStreamAdapterState::Uninit(stream),
        }
    }
}

impl RecordBatchStream for AsyncRecordBatchStreamAdapter {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn output_ordering(&self) -> Option<&[OrderOption]> {
        None
    }

    fn metrics(&self) -> Option<RecordBatchMetrics> {
        None
    }
}

impl Stream for AsyncRecordBatchStreamAdapter {
    type Item = Result<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            match &mut self.state {
                AsyncRecordBatchStreamAdapterState::Uninit(stream_future) => {
                    match ready!(Pin::new(stream_future).poll(cx)) {
                        Ok(stream) => {
                            self.state = AsyncRecordBatchStreamAdapterState::Ready(stream);
                            continue;
                        }
                        Err(e) => {
                            self.state = AsyncRecordBatchStreamAdapterState::Failed;
                            return Poll::Ready(Some(Err(e)));
                        }
                    };
                }
                AsyncRecordBatchStreamAdapterState::Ready(stream) => {
                    return Poll::Ready(ready!(Pin::new(stream).poll_next(cx)));
                }
                AsyncRecordBatchStreamAdapterState::Failed => return Poll::Ready(None),
            }
        }
    }

    // This is not supported for lazy stream.
    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, None)
    }
}

/// Custom cast function that handles Map -> Binary (JSON) conversion
fn custom_cast(
    array: &dyn Array,
    target_type: &ArrowDataType,
    extype: Option<ColumnExtType>,
) -> std::result::Result<Arc<dyn Array>, ArrowError> {
    if let ArrowDataType::Map(_, _) = array.data_type()
        && let ArrowDataType::Binary = target_type
    {
        return convert_map_to_json_binary(array, extype);
    }

    cast(array, target_type)
}

/// Convert a Map array to a Binary array containing JSON data
fn convert_map_to_json_binary(
    array: &dyn Array,
    extype: Option<ColumnExtType>,
) -> std::result::Result<Arc<dyn Array>, ArrowError> {
    use datatypes::arrow::array::{BinaryArray, MapArray};
    use serde_json::Value;

    let map_array = array
        .as_any()
        .downcast_ref::<MapArray>()
        .ok_or_else(|| ArrowError::CastError("Failed to downcast to MapArray".to_string()))?;

    let mut json_values = Vec::with_capacity(map_array.len());

    for i in 0..map_array.len() {
        if map_array.is_null(i) {
            json_values.push(None);
        } else {
            // Extract the map entry at index i
            let map_entry = map_array.value(i);
            let key_value_array = map_entry
                .as_any()
                .downcast_ref::<datatypes::arrow::array::StructArray>()
                .ok_or_else(|| {
                    ArrowError::CastError("Failed to downcast to StructArray".to_string())
                })?;

            // Convert to JSON object
            let mut json_obj = serde_json::Map::with_capacity(key_value_array.len());

            for j in 0..key_value_array.len() {
                if key_value_array.is_null(j) {
                    continue;
                }
                let key_field = key_value_array.column(0);
                let value_field = key_value_array.column(1);

                if key_field.is_null(j) {
                    continue;
                }

                let key = key_field
                    .as_any()
                    .downcast_ref::<datatypes::arrow::array::StringArray>()
                    .ok_or_else(|| {
                        ArrowError::CastError("Failed to downcast key to StringArray".to_string())
                    })?
                    .value(j);

                let value = if value_field.is_null(j) {
                    Value::Null
                } else {
                    let value_str = value_field
                        .as_any()
                        .downcast_ref::<datatypes::arrow::array::StringArray>()
                        .ok_or_else(|| {
                            ArrowError::CastError(
                                "Failed to downcast value to StringArray".to_string(),
                            )
                        })?
                        .value(j);
                    Value::String(value_str.to_string())
                };

                json_obj.insert(key.to_string(), value);
            }

            let json_value = Value::Object(json_obj);
            let json_bytes = match extype {
                Some(ColumnExtType::Json) => {
                    let json_string = match serde_json::to_string(&json_value) {
                        Ok(s) => s,
                        Err(e) => {
                            return Err(ArrowError::CastError(format!(
                                "Failed to serialize JSON: {}",
                                e
                            )));
                        }
                    };
                    match jsonb::parse_value(json_string.as_bytes()) {
                        Ok(jsonb_value) => jsonb_value.to_vec(),
                        Err(e) => {
                            return Err(ArrowError::CastError(format!(
                                "Failed to serialize JSONB: {}",
                                e
                            )));
                        }
                    }
                }
                _ => match serde_json::to_vec(&json_value) {
                    Ok(b) => b,
                    Err(e) => {
                        return Err(ArrowError::CastError(format!(
                            "Failed to serialize JSON: {}",
                            e
                        )));
                    }
                },
            };
            json_values.push(Some(json_bytes));
        }
    }

    let binary_array = BinaryArray::from_iter(json_values);
    Ok(Arc::new(binary_array))
}

#[cfg(test)]
mod test {
    use std::any::Any;
    use std::time::Duration;

    use common_error::ext::BoxedError;
    use common_error::mock::MockError;
    use common_error::status_code::StatusCode;
    use datafusion::execution::TaskContext;
    use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
    use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
    use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
    use datafusion::physical_plan::{DisplayAs, PlanProperties};
    use datatypes::arrow::array::{ArrayRef, MapArray, StringArray, StructArray};
    use datatypes::arrow::buffer::OffsetBuffer;
    use datatypes::arrow::datatypes::Field;
    use datatypes::prelude::ConcreteDataType;
    use datatypes::schema::ColumnSchema;
    use datatypes::vectors::Int32Vector;
    use futures::StreamExt;
    use serde_json::json;
    use snafu::IntoError;

    use super::*;
    use crate::RecordBatches;
    use crate::error::Error;

    #[derive(Debug)]
    struct TestMetricsExec {
        properties: Arc<PlanProperties>,
        metrics: ExecutionPlanMetricsSet,
    }

    impl TestMetricsExec {
        fn new(schema: DfSchemaRef) -> Self {
            Self::with_output_bytes(schema, &[24])
        }

        fn with_output_bytes(schema: DfSchemaRef, output_bytes_by_partition: &[usize]) -> Self {
            let metrics = ExecutionPlanMetricsSet::new();
            let elapsed_compute = MetricBuilder::new(&metrics).elapsed_compute(0);
            elapsed_compute.add_duration(Duration::from_nanos(42));
            for (partition, output_bytes) in output_bytes_by_partition.iter().copied().enumerate() {
                let metric = MetricBuilder::new(&metrics).output_bytes(partition);
                metric.add(output_bytes);
            }

            Self {
                properties: Arc::new(PlanProperties::new(
                    EquivalenceProperties::new(schema),
                    Partitioning::UnknownPartitioning(output_bytes_by_partition.len().max(1)),
                    EmissionType::Incremental,
                    Boundedness::Bounded,
                )),
                metrics,
            }
        }
    }

    impl DisplayAs for TestMetricsExec {
        fn fmt_as(&self, _t: DisplayFormatType, _f: &mut std::fmt::Formatter) -> std::fmt::Result {
            panic!("non-verbose lightweight partial metrics must not format the plan")
        }
    }

    impl ExecutionPlan for TestMetricsExec {
        fn name(&self) -> &str {
            REGION_SCAN_EXEC_NAME
        }

        fn as_any(&self) -> &dyn Any {
            self
        }

        fn properties(&self) -> &Arc<PlanProperties> {
            &self.properties
        }

        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
            vec![]
        }

        fn with_new_children(
            self: Arc<Self>,
            _children: Vec<Arc<dyn ExecutionPlan>>,
        ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
            Ok(self)
        }

        fn execute(
            &self,
            _partition: usize,
            _context: Arc<TaskContext>,
        ) -> datafusion_common::Result<DfSendableRecordBatchStream> {
            unreachable!("the test passes a separate stream to RecordBatchStreamAdapter")
        }

        fn metrics(&self) -> Option<MetricsSet> {
            Some(self.metrics.clone_inner())
        }
    }

    #[test]
    fn test_lightweight_query_load_metrics_sums_region_scan_output_bytes() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "a",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let plan = TestMetricsExec::with_output_bytes(schema.arrow_schema().clone(), &[24, 18]);

        let metrics = collect_lightweight_query_load_metrics(&plan, Some(42));

        assert_eq!(metrics.query_load_region_id, Some(42));
        assert_eq!(region_scan_output_bytes(&metrics), 42);
        assert_eq!(metrics.plan_metrics.len(), 1);
        assert_eq!(
            metrics.plan_metrics[0].metrics,
            vec![("output_bytes".to_string(), 42)]
        );
    }

    #[tokio::test]
    async fn test_record_batch_stream_adapter_collects_lightweight_partial_metrics() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "a",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let batch1 = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([1])) as _],
        )
        .unwrap()
        .into_df_record_batch();
        let batch2 = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([2])) as _],
        )
        .unwrap()
        .into_df_record_batch();
        let df_stream = Box::pin(
            datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
                schema.arrow_schema().clone(),
                futures::stream::iter(vec![Ok(batch1), Ok(batch2)]),
            ),
        );
        let plan = Arc::new(TestMetricsExec::new(schema.arrow_schema().clone()));

        let mut adapter = RecordBatchStreamAdapter::try_new(df_stream).unwrap();
        adapter.set_metrics2(plan);
        adapter.set_query_load_region_id(Some(42));

        assert!(adapter.metrics().is_none());
        assert!(adapter.next().await.unwrap().is_ok());
        let metrics = adapter
            .metrics()
            .expect("non-verbose queries need partial query-load metrics before EOF");
        assert_eq!(metrics.elapsed_compute, 42);
        assert_eq!(metrics.query_load_region_id, Some(42));
        assert_eq!(region_scan_output_bytes(&metrics), 24);
        assert_eq!(metrics.plan_metrics.len(), 1);
        assert_eq!(metrics.plan_metrics[0].plan, REGION_SCAN_EXEC_NAME);
        assert_eq!(metrics.plan_metrics[0].plan_name, REGION_SCAN_EXEC_NAME);
        assert_eq!(
            metrics.plan_metrics[0].metrics,
            vec![("output_bytes".to_string(), 24)]
        );
    }

    #[test]
    fn test_record_batch_stream_adapter_reuses_partial_query_stats_on_drop() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "a",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let df_stream = Box::pin(
            datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
                schema.arrow_schema().clone(),
                futures::stream::empty::<datafusion::error::Result<DfRecordBatch>>(),
            ),
        );
        let counters = RegionQueryStatCounters {
            query_cpu_time: Arc::new(AtomicU64::new(10)),
            query_scanned_bytes: Arc::new(AtomicU64::new(20)),
        };
        let stale_metrics = RecordBatchMetrics {
            elapsed_compute: 1,
            plan_metrics: vec![PlanMetrics {
                plan: REGION_SCAN_EXEC_NAME.to_string(),
                plan_name: REGION_SCAN_EXEC_NAME.to_string(),
                level: 0,
                metrics: vec![("output_bytes".to_string(), 2)],
            }],
            ..Default::default()
        };
        let adapter = RecordBatchStreamAdapter {
            schema: schema.clone(),
            stream: df_stream,
            metrics: None,
            metrics_2: Metrics::PartialResolved(
                Arc::new(TestMetricsExec::new(schema.arrow_schema().clone())),
                stale_metrics,
            ),
            query_load_region_id: None,
            query_stat_counters: Some(counters.clone()),
            explain_verbose: false,
            span: Span::current(),
        };

        drop(adapter);

        assert_eq!(counters.query_cpu_time.load(Ordering::Relaxed), 11);
        assert_eq!(counters.query_scanned_bytes.load(Ordering::Relaxed), 22);
    }

    #[tokio::test]
    async fn test_async_recordbatch_stream_adaptor() {
        struct MaybeErrorRecordBatchStream {
            items: Vec<Result<RecordBatch>>,
        }

        impl RecordBatchStream for MaybeErrorRecordBatchStream {
            fn schema(&self) -> SchemaRef {
                unimplemented!()
            }

            fn output_ordering(&self) -> Option<&[OrderOption]> {
                None
            }

            fn metrics(&self) -> Option<RecordBatchMetrics> {
                None
            }
        }

        impl Stream for MaybeErrorRecordBatchStream {
            type Item = Result<RecordBatch>;

            fn poll_next(
                mut self: Pin<&mut Self>,
                _: &mut Context<'_>,
            ) -> Poll<Option<Self::Item>> {
                if let Some(batch) = self.items.pop() {
                    Poll::Ready(Some(Ok(batch?)))
                } else {
                    Poll::Ready(None)
                }
            }
        }

        fn new_future_stream(
            maybe_recordbatches: Result<Vec<Result<RecordBatch>>>,
        ) -> FutureStream {
            Box::pin(async move {
                maybe_recordbatches
                    .map(|items| Box::pin(MaybeErrorRecordBatchStream { items }) as _)
            })
        }

        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "a",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let batch1 = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([1])) as _],
        )
        .unwrap();
        let batch2 = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([2])) as _],
        )
        .unwrap();

        let success_stream = new_future_stream(Ok(vec![Ok(batch1.clone()), Ok(batch2.clone())]));
        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), success_stream);
        let collected = RecordBatches::try_collect(Box::pin(adapter)).await.unwrap();
        assert_eq!(
            collected,
            RecordBatches::try_new(schema.clone(), vec![batch2.clone(), batch1.clone()]).unwrap()
        );

        let poll_err_stream = new_future_stream(Ok(vec![
            Ok(batch1.clone()),
            Err(error::ExternalSnafu
                .into_error(BoxedError::new(MockError::new(StatusCode::Unknown)))),
        ]));
        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), poll_err_stream);
        let err = RecordBatches::try_collect(Box::pin(adapter))
            .await
            .unwrap_err();
        assert!(
            matches!(err, Error::External { .. }),
            "unexpected err {err}"
        );

        let failed_to_init_stream =
            new_future_stream(Err(error::ExternalSnafu
                .into_error(BoxedError::new(MockError::new(StatusCode::Internal)))));
        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), failed_to_init_stream);
        let err = RecordBatches::try_collect(Box::pin(adapter))
            .await
            .unwrap_err();
        assert!(
            matches!(err, Error::External { .. }),
            "unexpected err {err}"
        );
    }

    #[test]
    fn test_convert_map_to_json_binary() {
        let keys = StringArray::from(vec![Some("a"), Some("b"), Some("c"), Some("x")]);
        let values = StringArray::from(vec![Some("1"), None, Some("3"), Some("42")]);
        let key_field = Arc::new(Field::new("key", ArrowDataType::Utf8, false));
        let value_field = Arc::new(Field::new("value", ArrowDataType::Utf8, true));
        let struct_type = ArrowDataType::Struct(vec![key_field, value_field].into());

        let entries_field = Arc::new(Field::new("entries", struct_type, false));

        let struct_array = StructArray::from(vec![
            (
                Arc::new(Field::new("key", ArrowDataType::Utf8, false)),
                Arc::new(keys) as ArrayRef,
            ),
            (
                Arc::new(Field::new("value", ArrowDataType::Utf8, true)),
                Arc::new(values) as ArrayRef,
            ),
        ]);

        let offsets = OffsetBuffer::from_lengths([3, 0, 1]);
        let nulls = datatypes::arrow::buffer::NullBuffer::from(vec![true, false, true]);

        let map_array = MapArray::new(
            entries_field,
            offsets,
            struct_array,
            Some(nulls), // nulls
            false,
        );

        let result = convert_map_to_json_binary(&map_array, None).unwrap();
        let binary_array = result
            .as_any()
            .downcast_ref::<datatypes::arrow::array::BinaryArray>()
            .unwrap();

        let expected_jsons = [
            Some(r#"{"a":"1","b":null,"c":"3"}"#),
            None,
            Some(r#"{"x":"42"}"#),
        ];

        for (i, _) in expected_jsons.iter().enumerate() {
            if let Some(expected) = &expected_jsons[i] {
                assert!(!binary_array.is_null(i));
                let actual_bytes = binary_array.value(i);
                let actual_str = std::str::from_utf8(actual_bytes).unwrap();
                assert_eq!(actual_str, *expected);
            } else {
                assert!(binary_array.is_null(i));
            }
        }

        let result_json =
            convert_map_to_json_binary(&map_array, Some(ColumnExtType::Json)).unwrap();
        let binary_array_json = result_json
            .as_any()
            .downcast_ref::<datatypes::arrow::array::BinaryArray>()
            .unwrap();

        for (i, _) in expected_jsons.iter().enumerate() {
            if expected_jsons[i].is_some() {
                assert!(!binary_array_json.is_null(i));
                let actual_bytes = binary_array_json.value(i);
                assert_ne!(actual_bytes, expected_jsons[i].unwrap().as_bytes());
            } else {
                assert!(binary_array_json.is_null(i));
            }
        }
    }

    #[test]
    fn test_record_query_stats_updates_region_counters() {
        let counters = RegionQueryStatCounters {
            query_cpu_time: Arc::new(AtomicU64::new(10)),
            query_scanned_bytes: Arc::new(AtomicU64::new(20)),
        };
        let metrics = RecordBatchMetrics {
            elapsed_compute: 2_000_000,
            plan_metrics: vec![PlanMetrics {
                plan: "RegionScanExec: region=1".to_string(),
                plan_name: REGION_SCAN_EXEC_NAME.to_string(),
                level: 0,
                metrics: vec![("output_bytes".to_string(), 42)],
            }],
            ..Default::default()
        };

        record_query_stats(&counters, &metrics);

        assert_eq!(counters.query_cpu_time.load(Ordering::Relaxed), 2_000_010);
        assert_eq!(counters.query_scanned_bytes.load(Ordering::Relaxed), 62);
    }

    #[test]
    fn test_record_batch_stream_adapter_records_query_stats_on_drop() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "a",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let df_stream = Box::pin(
            datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
                schema.arrow_schema().clone(),
                futures::stream::empty::<datafusion::error::Result<DfRecordBatch>>(),
            ),
        );
        let counters = RegionQueryStatCounters {
            query_cpu_time: Arc::new(AtomicU64::new(10)),
            query_scanned_bytes: Arc::new(AtomicU64::new(20)),
        };
        let metrics = RecordBatchMetrics {
            elapsed_compute: 2_000_000,
            plan_metrics: vec![PlanMetrics {
                plan: "RegionScanExec: region=1".to_string(),
                plan_name: REGION_SCAN_EXEC_NAME.to_string(),
                level: 0,
                metrics: vec![("output_bytes".to_string(), 42)],
            }],
            ..Default::default()
        };
        let adapter = RecordBatchStreamAdapter {
            schema,
            stream: df_stream,
            metrics: None,
            metrics_2: Metrics::Resolved(metrics),
            query_load_region_id: None,
            query_stat_counters: Some(counters.clone()),
            explain_verbose: false,
            span: Span::current(),
        };

        drop(adapter);

        assert_eq!(counters.query_cpu_time.load(Ordering::Relaxed), 2_000_010);
        assert_eq!(counters.query_scanned_bytes.load(Ordering::Relaxed), 62);
    }

    #[test]
    fn test_recordbatch_metrics_deserializes_without_region_watermarks() {
        let metrics: RecordBatchMetrics = serde_json::from_value(json!({
            "elapsed_compute": 12,
            "memory_usage": 34,
            "plan_metrics": []
        }))
        .unwrap();

        assert!(metrics.region_watermarks.is_empty());
        assert_eq!(metrics.elapsed_compute, 12);
        assert_eq!(metrics.memory_usage, 34);
    }

    #[test]
    fn test_plan_metrics_deserializes_without_plan_name() {
        let metrics: RecordBatchMetrics = serde_json::from_value(json!({
            "elapsed_compute": 12,
            "memory_usage": 34,
            "plan_metrics": [{
                "plan": "SeqScan: region=1",
                "level": 0,
                "metrics": []
            }]
        }))
        .unwrap();

        assert_eq!(metrics.plan_metrics[0].plan_name, "");
    }

    #[test]
    fn test_recordbatch_metrics_region_watermarks_serde_roundtrip() {
        let metrics = RecordBatchMetrics {
            region_watermarks: vec![
                RegionWatermarkEntry {
                    region_id: 1,
                    watermark: Some(100),
                },
                RegionWatermarkEntry {
                    region_id: 2,
                    watermark: None,
                },
            ],
            ..Default::default()
        };

        let value = serde_json::to_value(&metrics).unwrap();
        assert_eq!(
            value.get("region_watermarks").unwrap(),
            &json!([
                { "region_id": 1, "watermark": 100 },
                { "region_id": 2 }
            ])
        );

        let decoded: RecordBatchMetrics = serde_json::from_value(value).unwrap();
        assert_eq!(decoded.region_watermarks, metrics.region_watermarks);
    }

    #[test]
    fn test_recordbatch_metrics_skips_empty_region_watermarks_on_serialize() {
        let value = serde_json::to_value(RecordBatchMetrics::default()).unwrap();
        assert!(value.get("region_watermarks").is_none());
    }
}