common_recordbatch/
adapter.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::{self, Display};
16use std::future::Future;
17use std::marker::PhantomData;
18use std::pin::Pin;
19use std::str::FromStr;
20use std::sync::Arc;
21use std::task::{Context, Poll};
22
23use common_base::readable_size::ReadableSize;
24use common_telemetry::tracing::{Span, info_span};
25use common_time::util::format_nanoseconds_human_readable;
26use datafusion::arrow::compute::cast;
27use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
28use datafusion::error::Result as DfResult;
29use datafusion::execution::context::ExecutionProps;
30use datafusion::logical_expr::Expr;
31use datafusion::logical_expr::utils::conjunction;
32use datafusion::physical_expr::create_physical_expr;
33use datafusion::physical_plan::metrics::{BaselineMetrics, MetricValue};
34use datafusion::physical_plan::{
35    DisplayFormatType, ExecutionPlan, ExecutionPlanVisitor, PhysicalExpr,
36    RecordBatchStream as DfRecordBatchStream, accept,
37};
38use datafusion_common::arrow::error::ArrowError;
39use datafusion_common::{DataFusionError, ToDFSchema};
40use datatypes::arrow::array::Array;
41use datatypes::arrow::datatypes::DataType as ArrowDataType;
42use datatypes::schema::{ColumnExtType, Schema, SchemaRef};
43use futures::ready;
44use jsonb;
45use pin_project::pin_project;
46use snafu::ResultExt;
47
48use crate::error::{self, Result};
49use crate::filter::batch_filter;
50use crate::{
51    DfRecordBatch, DfSendableRecordBatchStream, OrderOption, RecordBatch, RecordBatchStream,
52    SendableRecordBatchStream, Stream,
53};
54
55type FutureStream =
56    Pin<Box<dyn std::future::Future<Output = Result<SendableRecordBatchStream>> + Send>>;
57
58/// Casts the `RecordBatch`es of `stream` against the `output_schema`.
59#[pin_project]
60pub struct RecordBatchStreamTypeAdapter<T, E> {
61    #[pin]
62    stream: T,
63    projected_schema: DfSchemaRef,
64    projection: Vec<usize>,
65    predicate: Option<Arc<dyn PhysicalExpr>>,
66    phantom: PhantomData<E>,
67}
68
69impl<T, E> RecordBatchStreamTypeAdapter<T, E>
70where
71    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
72    E: std::error::Error + Send + Sync + 'static,
73{
74    pub fn new(projected_schema: DfSchemaRef, stream: T, projection: Option<Vec<usize>>) -> Self {
75        let projection = if let Some(projection) = projection {
76            projection
77        } else {
78            (0..projected_schema.fields().len()).collect()
79        };
80
81        Self {
82            stream,
83            projected_schema,
84            projection,
85            predicate: None,
86            phantom: Default::default(),
87        }
88    }
89
90    pub fn with_filter(mut self, filters: Vec<Expr>) -> Result<Self> {
91        let filters = if let Some(expr) = conjunction(filters) {
92            let df_schema = self
93                .projected_schema
94                .clone()
95                .to_dfschema_ref()
96                .context(error::PhysicalExprSnafu)?;
97
98            let filters = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())
99                .context(error::PhysicalExprSnafu)?;
100            Some(filters)
101        } else {
102            None
103        };
104        self.predicate = filters;
105        Ok(self)
106    }
107}
108
109impl<T, E> DfRecordBatchStream for RecordBatchStreamTypeAdapter<T, E>
110where
111    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
112    E: std::error::Error + Send + Sync + 'static,
113{
114    fn schema(&self) -> DfSchemaRef {
115        self.projected_schema.clone()
116    }
117}
118
119impl<T, E> Stream for RecordBatchStreamTypeAdapter<T, E>
120where
121    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
122    E: std::error::Error + Send + Sync + 'static,
123{
124    type Item = DfResult<DfRecordBatch>;
125
126    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
127        let this = self.project();
128
129        let batch = futures::ready!(this.stream.poll_next(cx))
130            .map(|r| r.map_err(|e| DataFusionError::External(Box::new(e))));
131
132        let projected_schema = this.projected_schema.clone();
133        let projection = this.projection.clone();
134        let predicate = this.predicate.clone();
135
136        let batch = batch.map(|b| {
137            b.and_then(|b| {
138                let projected_column = b.project(&projection)?;
139                if projected_column.schema().fields.len() != projected_schema.fields.len() {
140                   return Err(DataFusionError::ArrowError(Box::new(ArrowError::SchemaError(format!(
141                        "Trying to cast a RecordBatch into an incompatible schema. RecordBatch: {}, Target: {}",
142                        projected_column.schema(),
143                        projected_schema,
144                    ))), None));
145                }
146
147                let mut columns = Vec::with_capacity(projected_schema.fields.len());
148                for (idx,field) in projected_schema.fields.iter().enumerate() {
149                    let column = projected_column.column(idx);
150                    let extype = field.metadata().get("greptime:type").and_then(|s| ColumnExtType::from_str(s).ok());
151                    let output = custom_cast(&column, field.data_type(), extype)?;
152                    columns.push(output)
153                }
154                let record_batch = DfRecordBatch::try_new(projected_schema, columns)?;
155                let record_batch = if let Some(predicate) = predicate {
156                    batch_filter(&record_batch, &predicate)?
157                } else {
158                    record_batch
159                };
160                Ok(record_batch)
161            })
162        });
163
164        Poll::Ready(batch)
165    }
166
167    #[inline]
168    fn size_hint(&self) -> (usize, Option<usize>) {
169        self.stream.size_hint()
170    }
171}
172
173/// Greptime SendableRecordBatchStream -> DataFusion RecordBatchStream.
174/// The reverse one is [RecordBatchStreamAdapter].
175pub struct DfRecordBatchStreamAdapter {
176    stream: SendableRecordBatchStream,
177}
178
179impl DfRecordBatchStreamAdapter {
180    pub fn new(stream: SendableRecordBatchStream) -> Self {
181        Self { stream }
182    }
183}
184
185impl DfRecordBatchStream for DfRecordBatchStreamAdapter {
186    fn schema(&self) -> DfSchemaRef {
187        self.stream.schema().arrow_schema().clone()
188    }
189}
190
191impl Stream for DfRecordBatchStreamAdapter {
192    type Item = DfResult<DfRecordBatch>;
193
194    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
195        match Pin::new(&mut self.stream).poll_next(cx) {
196            Poll::Pending => Poll::Pending,
197            Poll::Ready(Some(recordbatch)) => match recordbatch {
198                Ok(recordbatch) => Poll::Ready(Some(Ok(recordbatch.into_df_record_batch()))),
199                Err(e) => Poll::Ready(Some(Err(DataFusionError::External(Box::new(e))))),
200            },
201            Poll::Ready(None) => Poll::Ready(None),
202        }
203    }
204
205    #[inline]
206    fn size_hint(&self) -> (usize, Option<usize>) {
207        self.stream.size_hint()
208    }
209}
210
211/// DataFusion [SendableRecordBatchStream](DfSendableRecordBatchStream) -> Greptime [RecordBatchStream].
212/// The reverse one is [DfRecordBatchStreamAdapter].
213/// It can collect metrics from DataFusion execution plan.
214pub struct RecordBatchStreamAdapter {
215    schema: SchemaRef,
216    stream: DfSendableRecordBatchStream,
217    metrics: Option<BaselineMetrics>,
218    /// Aggregated plan-level metrics. Resolved after an [ExecutionPlan] is finished.
219    metrics_2: Metrics,
220    query_load_region_id: Option<u64>,
221    /// Display plan and metrics in verbose mode.
222    explain_verbose: bool,
223    span: Span,
224}
225
226/// Json encoded metrics. Contains metric from a whole plan tree.
227enum Metrics {
228    Unavailable,
229    Unresolved(Arc<dyn ExecutionPlan>),
230    PartialResolved(Arc<dyn ExecutionPlan>, RecordBatchMetrics),
231    Resolved(RecordBatchMetrics),
232}
233
234impl RecordBatchStreamAdapter {
235    pub fn try_new(stream: DfSendableRecordBatchStream) -> Result<Self> {
236        let schema =
237            Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
238        Ok(Self {
239            schema,
240            stream,
241            metrics: None,
242            metrics_2: Metrics::Unavailable,
243            query_load_region_id: None,
244            explain_verbose: false,
245            span: Span::current(),
246        })
247    }
248
249    pub fn try_new_with_span(stream: DfSendableRecordBatchStream, span: Span) -> Result<Self> {
250        let schema =
251            Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
252        let subspan = info_span!(parent: &span, "RecordBatchStreamAdapter");
253        Ok(Self {
254            schema,
255            stream,
256            metrics: None,
257            metrics_2: Metrics::Unavailable,
258            query_load_region_id: None,
259            explain_verbose: false,
260            span: subspan,
261        })
262    }
263
264    pub fn set_metrics2(&mut self, plan: Arc<dyn ExecutionPlan>) {
265        self.metrics_2 = Metrics::Unresolved(plan)
266    }
267
268    pub fn set_query_load_region_id(&mut self, region_id: Option<u64>) {
269        self.query_load_region_id = region_id;
270    }
271
272    /// Set the verbose mode for displaying plan and metrics.
273    pub fn set_explain_verbose(&mut self, verbose: bool) {
274        self.explain_verbose = verbose;
275    }
276}
277
278impl RecordBatchStream for RecordBatchStreamAdapter {
279    fn name(&self) -> &str {
280        "RecordBatchStreamAdapter"
281    }
282
283    fn schema(&self) -> SchemaRef {
284        self.schema.clone()
285    }
286
287    fn metrics(&self) -> Option<RecordBatchMetrics> {
288        match &self.metrics_2 {
289            Metrics::Resolved(metrics) | Metrics::PartialResolved(_, metrics) => {
290                Some(metrics.clone())
291            }
292            Metrics::Unavailable | Metrics::Unresolved(_) => None,
293        }
294    }
295
296    fn output_ordering(&self) -> Option<&[OrderOption]> {
297        None
298    }
299}
300
301impl Stream for RecordBatchStreamAdapter {
302    type Item = Result<RecordBatch>;
303
304    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
305        let timer = self
306            .metrics
307            .as_ref()
308            .map(|m| m.elapsed_compute().clone())
309            .unwrap_or_default();
310        let _guard = timer.timer();
311        let poll_span = info_span!(parent: &self.span, "poll_next");
312        let _entered = poll_span.enter();
313        match Pin::new(&mut self.stream).poll_next(cx) {
314            Poll::Pending => Poll::Pending,
315            Poll::Ready(Some(df_record_batch)) => {
316                let df_record_batch = df_record_batch?;
317                if let Metrics::Unresolved(df_plan) | Metrics::PartialResolved(df_plan, _) =
318                    &self.metrics_2
319                {
320                    let mut metric_collector = MetricCollector::new(self.explain_verbose);
321                    accept(df_plan.as_ref(), &mut metric_collector).unwrap();
322                    metric_collector.record_batch_metrics.query_load_region_id =
323                        self.query_load_region_id;
324                    self.metrics_2 = Metrics::PartialResolved(
325                        df_plan.clone(),
326                        metric_collector.record_batch_metrics,
327                    );
328                }
329                Poll::Ready(Some(Ok(RecordBatch::from_df_record_batch(
330                    self.schema(),
331                    df_record_batch,
332                ))))
333            }
334            Poll::Ready(None) => {
335                if let Metrics::Unresolved(df_plan) | Metrics::PartialResolved(df_plan, _) =
336                    &self.metrics_2
337                {
338                    let mut metric_collector = MetricCollector::new(self.explain_verbose);
339                    accept(df_plan.as_ref(), &mut metric_collector).unwrap();
340                    metric_collector.record_batch_metrics.query_load_region_id =
341                        self.query_load_region_id;
342                    self.metrics_2 = Metrics::Resolved(metric_collector.record_batch_metrics);
343                }
344                Poll::Ready(None)
345            }
346        }
347    }
348
349    #[inline]
350    fn size_hint(&self) -> (usize, Option<usize>) {
351        self.stream.size_hint()
352    }
353}
354
355/// An [ExecutionPlanVisitor] to collect metrics from a [ExecutionPlan].
356pub struct MetricCollector {
357    current_level: usize,
358    pub record_batch_metrics: RecordBatchMetrics,
359    verbose: bool,
360}
361
362impl MetricCollector {
363    pub fn new(verbose: bool) -> Self {
364        Self {
365            current_level: 0,
366            record_batch_metrics: RecordBatchMetrics::default(),
367            verbose,
368        }
369    }
370}
371
372impl ExecutionPlanVisitor for MetricCollector {
373    type Error = !;
374
375    fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
376        // skip if no metric available
377        let Some(metric) = plan.metrics() else {
378            self.record_batch_metrics.plan_metrics.push(PlanMetrics {
379                plan: plan.name().to_string(),
380                plan_name: plan.name().to_string(),
381                level: self.current_level,
382                metrics: vec![],
383            });
384            self.current_level += 1;
385            return Ok(true);
386        };
387
388        // scrape plan metrics
389        let metric = metric
390            .aggregate_by_name()
391            .sorted_for_display()
392            .timestamps_removed();
393        let mut plan_metric = PlanMetrics {
394            plan: one_line(plan, self.verbose).to_string(),
395            plan_name: plan.name().to_string(),
396            level: self.current_level,
397            metrics: Vec::with_capacity(metric.iter().size_hint().0),
398        };
399        for m in metric.iter() {
400            plan_metric
401                .metrics
402                .push((m.value().name().to_string(), m.value().as_usize()));
403
404            // aggregate high-level metrics
405            match m.value() {
406                MetricValue::ElapsedCompute(ec) => {
407                    self.record_batch_metrics.elapsed_compute += ec.value()
408                }
409                MetricValue::CurrentMemoryUsage(m) => {
410                    self.record_batch_metrics.memory_usage += m.value()
411                }
412                _ => {}
413            }
414        }
415        self.record_batch_metrics.plan_metrics.push(plan_metric);
416
417        self.current_level += 1;
418        Ok(true)
419    }
420
421    fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
422        self.current_level -= 1;
423        Ok(true)
424    }
425}
426
427/// Returns a single-line summary of the root of the plan.
428/// If the `verbose` flag is set, it will display detailed information about the plan.
429fn one_line(plan: &dyn ExecutionPlan, verbose: bool) -> impl fmt::Display + '_ {
430    struct Wrapper<'a> {
431        plan: &'a dyn ExecutionPlan,
432        format_type: DisplayFormatType,
433    }
434
435    impl fmt::Display for Wrapper<'_> {
436        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
437            self.plan.fmt_as(self.format_type, f)?;
438            writeln!(f)
439        }
440    }
441
442    let format_type = if verbose {
443        DisplayFormatType::Verbose
444    } else {
445        DisplayFormatType::Default
446    };
447    Wrapper { plan, format_type }
448}
449
450/// [`RecordBatchMetrics`] carrys metrics value
451/// from datanode to frontend through gRPC
452#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
453pub struct RecordBatchMetrics {
454    // High-level aggregated metrics
455    /// CPU consumption in nanoseconds
456    pub elapsed_compute: usize,
457    /// Memory used by the plan in bytes
458    pub memory_usage: usize,
459    // Detailed per-plan metrics
460    /// An ordered list of plan metrics, from top to bottom in post-order.
461    pub plan_metrics: Vec<PlanMetrics>,
462    /// Region id that should receive query-load metrics for this scan.
463    #[serde(default, skip_serializing_if = "Option::is_none")]
464    pub query_load_region_id: Option<u64>,
465    /// Per-region watermark for incremental-read checkpoint advancement.
466    ///
467    /// The watermark is the latest sequence (`seq`) this query round safely read
468    /// for each participating region. Flow uses it to decide where the next
469    /// incremental round can resume.
470    ///
471    /// - `Some(seq)`: the query proved it safely read up to `seq`; downstream
472    ///   may advance the checkpoint to this value.
473    /// - `None`: the region participated but the query could not prove a safe
474    ///   read upper-bound, so the checkpoint must not advance for this region.
475    ///
476    /// Omitted when empty for backward compatibility.
477    #[serde(default, skip_serializing_if = "Vec::is_empty")]
478    pub region_watermarks: Vec<RegionWatermarkEntry>,
479}
480
481#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
482pub struct RegionWatermarkEntry {
483    pub region_id: u64,
484    #[serde(default, skip_serializing_if = "Option::is_none")]
485    pub watermark: Option<u64>,
486}
487
488/// Determines if a metric name represents a time measurement that should be formatted.
489fn is_time_metric(metric_name: &str) -> bool {
490    metric_name.contains("elapsed") || metric_name.contains("time") || metric_name.contains("cost")
491}
492
493/// Determines if a metric name represents a bytes measurement that should be formatted.
494fn is_bytes_metric(metric_name: &str) -> bool {
495    metric_name.contains("bytes") || metric_name.contains("mem")
496}
497
498fn format_bytes_human_readable(bytes: usize) -> String {
499    format!("{}", ReadableSize(bytes as u64))
500}
501
502/// Only display `plan_metrics` with indent `  ` (2 spaces).
503impl Display for RecordBatchMetrics {
504    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
505        for metric in &self.plan_metrics {
506            write!(
507                f,
508                "{:indent$}{} metrics=[",
509                " ",
510                metric.plan.trim_end(),
511                indent = metric.level * 2,
512            )?;
513            for (label, value) in &metric.metrics {
514                if is_time_metric(label) {
515                    write!(
516                        f,
517                        "{}: {}, ",
518                        label,
519                        format_nanoseconds_human_readable(*value),
520                    )?;
521                } else if is_bytes_metric(label) {
522                    write!(f, "{}: {}, ", label, format_bytes_human_readable(*value),)?;
523                } else {
524                    write!(f, "{}: {}, ", label, value)?;
525                }
526            }
527            writeln!(f, "]")?;
528        }
529
530        Ok(())
531    }
532}
533
534#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
535pub struct PlanMetrics {
536    /// The plan name
537    pub plan: String,
538    /// The stable execution plan name.
539    #[serde(default)]
540    pub plan_name: String,
541    /// The level of the plan, starts from 0
542    pub level: usize,
543    /// An ordered key-value list of metrics.
544    /// Key is metric label and value is metric value.
545    pub metrics: Vec<(String, usize)>,
546}
547
548enum AsyncRecordBatchStreamAdapterState {
549    Uninit(FutureStream),
550    Ready(SendableRecordBatchStream),
551    Failed,
552}
553
554pub struct AsyncRecordBatchStreamAdapter {
555    schema: SchemaRef,
556    state: AsyncRecordBatchStreamAdapterState,
557}
558
559impl AsyncRecordBatchStreamAdapter {
560    pub fn new(schema: SchemaRef, stream: FutureStream) -> Self {
561        Self {
562            schema,
563            state: AsyncRecordBatchStreamAdapterState::Uninit(stream),
564        }
565    }
566}
567
568impl RecordBatchStream for AsyncRecordBatchStreamAdapter {
569    fn schema(&self) -> SchemaRef {
570        self.schema.clone()
571    }
572
573    fn output_ordering(&self) -> Option<&[OrderOption]> {
574        None
575    }
576
577    fn metrics(&self) -> Option<RecordBatchMetrics> {
578        None
579    }
580}
581
582impl Stream for AsyncRecordBatchStreamAdapter {
583    type Item = Result<RecordBatch>;
584
585    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
586        loop {
587            match &mut self.state {
588                AsyncRecordBatchStreamAdapterState::Uninit(stream_future) => {
589                    match ready!(Pin::new(stream_future).poll(cx)) {
590                        Ok(stream) => {
591                            self.state = AsyncRecordBatchStreamAdapterState::Ready(stream);
592                            continue;
593                        }
594                        Err(e) => {
595                            self.state = AsyncRecordBatchStreamAdapterState::Failed;
596                            return Poll::Ready(Some(Err(e)));
597                        }
598                    };
599                }
600                AsyncRecordBatchStreamAdapterState::Ready(stream) => {
601                    return Poll::Ready(ready!(Pin::new(stream).poll_next(cx)));
602                }
603                AsyncRecordBatchStreamAdapterState::Failed => return Poll::Ready(None),
604            }
605        }
606    }
607
608    // This is not supported for lazy stream.
609    #[inline]
610    fn size_hint(&self) -> (usize, Option<usize>) {
611        (0, None)
612    }
613}
614
615/// Custom cast function that handles Map -> Binary (JSON) conversion
616fn custom_cast(
617    array: &dyn Array,
618    target_type: &ArrowDataType,
619    extype: Option<ColumnExtType>,
620) -> std::result::Result<Arc<dyn Array>, ArrowError> {
621    if let ArrowDataType::Map(_, _) = array.data_type()
622        && let ArrowDataType::Binary = target_type
623    {
624        return convert_map_to_json_binary(array, extype);
625    }
626
627    cast(array, target_type)
628}
629
630/// Convert a Map array to a Binary array containing JSON data
631fn convert_map_to_json_binary(
632    array: &dyn Array,
633    extype: Option<ColumnExtType>,
634) -> std::result::Result<Arc<dyn Array>, ArrowError> {
635    use datatypes::arrow::array::{BinaryArray, MapArray};
636    use serde_json::Value;
637
638    let map_array = array
639        .as_any()
640        .downcast_ref::<MapArray>()
641        .ok_or_else(|| ArrowError::CastError("Failed to downcast to MapArray".to_string()))?;
642
643    let mut json_values = Vec::with_capacity(map_array.len());
644
645    for i in 0..map_array.len() {
646        if map_array.is_null(i) {
647            json_values.push(None);
648        } else {
649            // Extract the map entry at index i
650            let map_entry = map_array.value(i);
651            let key_value_array = map_entry
652                .as_any()
653                .downcast_ref::<datatypes::arrow::array::StructArray>()
654                .ok_or_else(|| {
655                    ArrowError::CastError("Failed to downcast to StructArray".to_string())
656                })?;
657
658            // Convert to JSON object
659            let mut json_obj = serde_json::Map::with_capacity(key_value_array.len());
660
661            for j in 0..key_value_array.len() {
662                if key_value_array.is_null(j) {
663                    continue;
664                }
665                let key_field = key_value_array.column(0);
666                let value_field = key_value_array.column(1);
667
668                if key_field.is_null(j) {
669                    continue;
670                }
671
672                let key = key_field
673                    .as_any()
674                    .downcast_ref::<datatypes::arrow::array::StringArray>()
675                    .ok_or_else(|| {
676                        ArrowError::CastError("Failed to downcast key to StringArray".to_string())
677                    })?
678                    .value(j);
679
680                let value = if value_field.is_null(j) {
681                    Value::Null
682                } else {
683                    let value_str = value_field
684                        .as_any()
685                        .downcast_ref::<datatypes::arrow::array::StringArray>()
686                        .ok_or_else(|| {
687                            ArrowError::CastError(
688                                "Failed to downcast value to StringArray".to_string(),
689                            )
690                        })?
691                        .value(j);
692                    Value::String(value_str.to_string())
693                };
694
695                json_obj.insert(key.to_string(), value);
696            }
697
698            let json_value = Value::Object(json_obj);
699            let json_bytes = match extype {
700                Some(ColumnExtType::Json) => {
701                    let json_string = match serde_json::to_string(&json_value) {
702                        Ok(s) => s,
703                        Err(e) => {
704                            return Err(ArrowError::CastError(format!(
705                                "Failed to serialize JSON: {}",
706                                e
707                            )));
708                        }
709                    };
710                    match jsonb::parse_value(json_string.as_bytes()) {
711                        Ok(jsonb_value) => jsonb_value.to_vec(),
712                        Err(e) => {
713                            return Err(ArrowError::CastError(format!(
714                                "Failed to serialize JSONB: {}",
715                                e
716                            )));
717                        }
718                    }
719                }
720                _ => match serde_json::to_vec(&json_value) {
721                    Ok(b) => b,
722                    Err(e) => {
723                        return Err(ArrowError::CastError(format!(
724                            "Failed to serialize JSON: {}",
725                            e
726                        )));
727                    }
728                },
729            };
730            json_values.push(Some(json_bytes));
731        }
732    }
733
734    let binary_array = BinaryArray::from_iter(json_values);
735    Ok(Arc::new(binary_array))
736}
737
738#[cfg(test)]
739mod test {
740    use common_error::ext::BoxedError;
741    use common_error::mock::MockError;
742    use common_error::status_code::StatusCode;
743    use datatypes::arrow::array::{ArrayRef, MapArray, StringArray, StructArray};
744    use datatypes::arrow::buffer::OffsetBuffer;
745    use datatypes::arrow::datatypes::Field;
746    use datatypes::prelude::ConcreteDataType;
747    use datatypes::schema::ColumnSchema;
748    use datatypes::vectors::Int32Vector;
749    use serde_json::json;
750    use snafu::IntoError;
751
752    use super::*;
753    use crate::RecordBatches;
754    use crate::error::Error;
755
756    #[tokio::test]
757    async fn test_async_recordbatch_stream_adaptor() {
758        struct MaybeErrorRecordBatchStream {
759            items: Vec<Result<RecordBatch>>,
760        }
761
762        impl RecordBatchStream for MaybeErrorRecordBatchStream {
763            fn schema(&self) -> SchemaRef {
764                unimplemented!()
765            }
766
767            fn output_ordering(&self) -> Option<&[OrderOption]> {
768                None
769            }
770
771            fn metrics(&self) -> Option<RecordBatchMetrics> {
772                None
773            }
774        }
775
776        impl Stream for MaybeErrorRecordBatchStream {
777            type Item = Result<RecordBatch>;
778
779            fn poll_next(
780                mut self: Pin<&mut Self>,
781                _: &mut Context<'_>,
782            ) -> Poll<Option<Self::Item>> {
783                if let Some(batch) = self.items.pop() {
784                    Poll::Ready(Some(Ok(batch?)))
785                } else {
786                    Poll::Ready(None)
787                }
788            }
789        }
790
791        fn new_future_stream(
792            maybe_recordbatches: Result<Vec<Result<RecordBatch>>>,
793        ) -> FutureStream {
794            Box::pin(async move {
795                maybe_recordbatches
796                    .map(|items| Box::pin(MaybeErrorRecordBatchStream { items }) as _)
797            })
798        }
799
800        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
801            "a",
802            ConcreteDataType::int32_datatype(),
803            false,
804        )]));
805        let batch1 = RecordBatch::new(
806            schema.clone(),
807            vec![Arc::new(Int32Vector::from_slice([1])) as _],
808        )
809        .unwrap();
810        let batch2 = RecordBatch::new(
811            schema.clone(),
812            vec![Arc::new(Int32Vector::from_slice([2])) as _],
813        )
814        .unwrap();
815
816        let success_stream = new_future_stream(Ok(vec![Ok(batch1.clone()), Ok(batch2.clone())]));
817        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), success_stream);
818        let collected = RecordBatches::try_collect(Box::pin(adapter)).await.unwrap();
819        assert_eq!(
820            collected,
821            RecordBatches::try_new(schema.clone(), vec![batch2.clone(), batch1.clone()]).unwrap()
822        );
823
824        let poll_err_stream = new_future_stream(Ok(vec![
825            Ok(batch1.clone()),
826            Err(error::ExternalSnafu
827                .into_error(BoxedError::new(MockError::new(StatusCode::Unknown)))),
828        ]));
829        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), poll_err_stream);
830        let err = RecordBatches::try_collect(Box::pin(adapter))
831            .await
832            .unwrap_err();
833        assert!(
834            matches!(err, Error::External { .. }),
835            "unexpected err {err}"
836        );
837
838        let failed_to_init_stream =
839            new_future_stream(Err(error::ExternalSnafu
840                .into_error(BoxedError::new(MockError::new(StatusCode::Internal)))));
841        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), failed_to_init_stream);
842        let err = RecordBatches::try_collect(Box::pin(adapter))
843            .await
844            .unwrap_err();
845        assert!(
846            matches!(err, Error::External { .. }),
847            "unexpected err {err}"
848        );
849    }
850
851    #[test]
852    fn test_convert_map_to_json_binary() {
853        let keys = StringArray::from(vec![Some("a"), Some("b"), Some("c"), Some("x")]);
854        let values = StringArray::from(vec![Some("1"), None, Some("3"), Some("42")]);
855        let key_field = Arc::new(Field::new("key", ArrowDataType::Utf8, false));
856        let value_field = Arc::new(Field::new("value", ArrowDataType::Utf8, true));
857        let struct_type = ArrowDataType::Struct(vec![key_field, value_field].into());
858
859        let entries_field = Arc::new(Field::new("entries", struct_type, false));
860
861        let struct_array = StructArray::from(vec![
862            (
863                Arc::new(Field::new("key", ArrowDataType::Utf8, false)),
864                Arc::new(keys) as ArrayRef,
865            ),
866            (
867                Arc::new(Field::new("value", ArrowDataType::Utf8, true)),
868                Arc::new(values) as ArrayRef,
869            ),
870        ]);
871
872        let offsets = OffsetBuffer::from_lengths([3, 0, 1]);
873        let nulls = datatypes::arrow::buffer::NullBuffer::from(vec![true, false, true]);
874
875        let map_array = MapArray::new(
876            entries_field,
877            offsets,
878            struct_array,
879            Some(nulls), // nulls
880            false,
881        );
882
883        let result = convert_map_to_json_binary(&map_array, None).unwrap();
884        let binary_array = result
885            .as_any()
886            .downcast_ref::<datatypes::arrow::array::BinaryArray>()
887            .unwrap();
888
889        let expected_jsons = [
890            Some(r#"{"a":"1","b":null,"c":"3"}"#),
891            None,
892            Some(r#"{"x":"42"}"#),
893        ];
894
895        for (i, _) in expected_jsons.iter().enumerate() {
896            if let Some(expected) = &expected_jsons[i] {
897                assert!(!binary_array.is_null(i));
898                let actual_bytes = binary_array.value(i);
899                let actual_str = std::str::from_utf8(actual_bytes).unwrap();
900                assert_eq!(actual_str, *expected);
901            } else {
902                assert!(binary_array.is_null(i));
903            }
904        }
905
906        let result_json =
907            convert_map_to_json_binary(&map_array, Some(ColumnExtType::Json)).unwrap();
908        let binary_array_json = result_json
909            .as_any()
910            .downcast_ref::<datatypes::arrow::array::BinaryArray>()
911            .unwrap();
912
913        for (i, _) in expected_jsons.iter().enumerate() {
914            if expected_jsons[i].is_some() {
915                assert!(!binary_array_json.is_null(i));
916                let actual_bytes = binary_array_json.value(i);
917                assert_ne!(actual_bytes, expected_jsons[i].unwrap().as_bytes());
918            } else {
919                assert!(binary_array_json.is_null(i));
920            }
921        }
922    }
923
924    #[test]
925    fn test_recordbatch_metrics_deserializes_without_region_watermarks() {
926        let metrics: RecordBatchMetrics = serde_json::from_value(json!({
927            "elapsed_compute": 12,
928            "memory_usage": 34,
929            "plan_metrics": []
930        }))
931        .unwrap();
932
933        assert!(metrics.region_watermarks.is_empty());
934        assert_eq!(metrics.elapsed_compute, 12);
935        assert_eq!(metrics.memory_usage, 34);
936    }
937
938    #[test]
939    fn test_plan_metrics_deserializes_without_plan_name() {
940        let metrics: RecordBatchMetrics = serde_json::from_value(json!({
941            "elapsed_compute": 12,
942            "memory_usage": 34,
943            "plan_metrics": [{
944                "plan": "SeqScan: region=1",
945                "level": 0,
946                "metrics": []
947            }]
948        }))
949        .unwrap();
950
951        assert_eq!(metrics.plan_metrics[0].plan_name, "");
952    }
953
954    #[test]
955    fn test_recordbatch_metrics_region_watermarks_serde_roundtrip() {
956        let metrics = RecordBatchMetrics {
957            region_watermarks: vec![
958                RegionWatermarkEntry {
959                    region_id: 1,
960                    watermark: Some(100),
961                },
962                RegionWatermarkEntry {
963                    region_id: 2,
964                    watermark: None,
965                },
966            ],
967            ..Default::default()
968        };
969
970        let value = serde_json::to_value(&metrics).unwrap();
971        assert_eq!(
972            value.get("region_watermarks").unwrap(),
973            &json!([
974                { "region_id": 1, "watermark": 100 },
975                { "region_id": 2 }
976            ])
977        );
978
979        let decoded: RecordBatchMetrics = serde_json::from_value(value).unwrap();
980        assert_eq!(decoded.region_watermarks, metrics.region_watermarks);
981    }
982
983    #[test]
984    fn test_recordbatch_metrics_skips_empty_region_watermarks_on_serialize() {
985        let value = serde_json::to_value(RecordBatchMetrics::default()).unwrap();
986        assert!(value.get("region_watermarks").is_none());
987    }
988}
common_recordbatch/adapter.rs

common_recordbatch/
adapter.rs
