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    /// Display plan and metrics in verbose mode.
221    explain_verbose: bool,
222    span: Span,
223}
224
225/// Json encoded metrics. Contains metric from a whole plan tree.
226enum Metrics {
227    Unavailable,
228    Unresolved(Arc<dyn ExecutionPlan>),
229    PartialResolved(Arc<dyn ExecutionPlan>, RecordBatchMetrics),
230    Resolved(RecordBatchMetrics),
231}
232
233impl RecordBatchStreamAdapter {
234    pub fn try_new(stream: DfSendableRecordBatchStream) -> Result<Self> {
235        let schema =
236            Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
237        Ok(Self {
238            schema,
239            stream,
240            metrics: None,
241            metrics_2: Metrics::Unavailable,
242            explain_verbose: false,
243            span: Span::current(),
244        })
245    }
246
247    pub fn try_new_with_span(stream: DfSendableRecordBatchStream, span: Span) -> Result<Self> {
248        let schema =
249            Arc::new(Schema::try_from(stream.schema()).context(error::SchemaConversionSnafu)?);
250        let subspan = info_span!(parent: &span, "RecordBatchStreamAdapter");
251        Ok(Self {
252            schema,
253            stream,
254            metrics: None,
255            metrics_2: Metrics::Unavailable,
256            explain_verbose: false,
257            span: subspan,
258        })
259    }
260
261    pub fn set_metrics2(&mut self, plan: Arc<dyn ExecutionPlan>) {
262        self.metrics_2 = Metrics::Unresolved(plan)
263    }
264
265    /// Set the verbose mode for displaying plan and metrics.
266    pub fn set_explain_verbose(&mut self, verbose: bool) {
267        self.explain_verbose = verbose;
268    }
269}
270
271impl RecordBatchStream for RecordBatchStreamAdapter {
272    fn name(&self) -> &str {
273        "RecordBatchStreamAdapter"
274    }
275
276    fn schema(&self) -> SchemaRef {
277        self.schema.clone()
278    }
279
280    fn metrics(&self) -> Option<RecordBatchMetrics> {
281        match &self.metrics_2 {
282            Metrics::Resolved(metrics) | Metrics::PartialResolved(_, metrics) => {
283                Some(metrics.clone())
284            }
285            Metrics::Unavailable | Metrics::Unresolved(_) => None,
286        }
287    }
288
289    fn output_ordering(&self) -> Option<&[OrderOption]> {
290        None
291    }
292}
293
294impl Stream for RecordBatchStreamAdapter {
295    type Item = Result<RecordBatch>;
296
297    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
298        let timer = self
299            .metrics
300            .as_ref()
301            .map(|m| m.elapsed_compute().clone())
302            .unwrap_or_default();
303        let _guard = timer.timer();
304        let poll_span = info_span!(parent: &self.span, "poll_next");
305        let _entered = poll_span.enter();
306        match Pin::new(&mut self.stream).poll_next(cx) {
307            Poll::Pending => Poll::Pending,
308            Poll::Ready(Some(df_record_batch)) => {
309                let df_record_batch = df_record_batch?;
310                if let Metrics::Unresolved(df_plan) | Metrics::PartialResolved(df_plan, _) =
311                    &self.metrics_2
312                {
313                    let mut metric_collector = MetricCollector::new(self.explain_verbose);
314                    accept(df_plan.as_ref(), &mut metric_collector).unwrap();
315                    self.metrics_2 = Metrics::PartialResolved(
316                        df_plan.clone(),
317                        metric_collector.record_batch_metrics,
318                    );
319                }
320                Poll::Ready(Some(Ok(RecordBatch::from_df_record_batch(
321                    self.schema(),
322                    df_record_batch,
323                ))))
324            }
325            Poll::Ready(None) => {
326                if let Metrics::Unresolved(df_plan) | Metrics::PartialResolved(df_plan, _) =
327                    &self.metrics_2
328                {
329                    let mut metric_collector = MetricCollector::new(self.explain_verbose);
330                    accept(df_plan.as_ref(), &mut metric_collector).unwrap();
331                    self.metrics_2 = Metrics::Resolved(metric_collector.record_batch_metrics);
332                }
333                Poll::Ready(None)
334            }
335        }
336    }
337
338    #[inline]
339    fn size_hint(&self) -> (usize, Option<usize>) {
340        self.stream.size_hint()
341    }
342}
343
344/// An [ExecutionPlanVisitor] to collect metrics from a [ExecutionPlan].
345pub struct MetricCollector {
346    current_level: usize,
347    pub record_batch_metrics: RecordBatchMetrics,
348    verbose: bool,
349}
350
351impl MetricCollector {
352    pub fn new(verbose: bool) -> Self {
353        Self {
354            current_level: 0,
355            record_batch_metrics: RecordBatchMetrics::default(),
356            verbose,
357        }
358    }
359}
360
361impl ExecutionPlanVisitor for MetricCollector {
362    type Error = !;
363
364    fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
365        // skip if no metric available
366        let Some(metric) = plan.metrics() else {
367            self.record_batch_metrics.plan_metrics.push(PlanMetrics {
368                plan: plan.name().to_string(),
369                plan_name: plan.name().to_string(),
370                level: self.current_level,
371                metrics: vec![],
372            });
373            self.current_level += 1;
374            return Ok(true);
375        };
376
377        // scrape plan metrics
378        let metric = metric
379            .aggregate_by_name()
380            .sorted_for_display()
381            .timestamps_removed();
382        let mut plan_metric = PlanMetrics {
383            plan: one_line(plan, self.verbose).to_string(),
384            plan_name: plan.name().to_string(),
385            level: self.current_level,
386            metrics: Vec::with_capacity(metric.iter().size_hint().0),
387        };
388        for m in metric.iter() {
389            plan_metric
390                .metrics
391                .push((m.value().name().to_string(), m.value().as_usize()));
392
393            // aggregate high-level metrics
394            match m.value() {
395                MetricValue::ElapsedCompute(ec) => {
396                    self.record_batch_metrics.elapsed_compute += ec.value()
397                }
398                MetricValue::CurrentMemoryUsage(m) => {
399                    self.record_batch_metrics.memory_usage += m.value()
400                }
401                _ => {}
402            }
403        }
404        self.record_batch_metrics.plan_metrics.push(plan_metric);
405
406        self.current_level += 1;
407        Ok(true)
408    }
409
410    fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> std::result::Result<bool, Self::Error> {
411        self.current_level -= 1;
412        Ok(true)
413    }
414}
415
416/// Returns a single-line summary of the root of the plan.
417/// If the `verbose` flag is set, it will display detailed information about the plan.
418fn one_line(plan: &dyn ExecutionPlan, verbose: bool) -> impl fmt::Display + '_ {
419    struct Wrapper<'a> {
420        plan: &'a dyn ExecutionPlan,
421        format_type: DisplayFormatType,
422    }
423
424    impl fmt::Display for Wrapper<'_> {
425        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
426            self.plan.fmt_as(self.format_type, f)?;
427            writeln!(f)
428        }
429    }
430
431    let format_type = if verbose {
432        DisplayFormatType::Verbose
433    } else {
434        DisplayFormatType::Default
435    };
436    Wrapper { plan, format_type }
437}
438
439/// [`RecordBatchMetrics`] carrys metrics value
440/// from datanode to frontend through gRPC
441#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
442pub struct RecordBatchMetrics {
443    // High-level aggregated metrics
444    /// CPU consumption in nanoseconds
445    pub elapsed_compute: usize,
446    /// Memory used by the plan in bytes
447    pub memory_usage: usize,
448    // Detailed per-plan metrics
449    /// An ordered list of plan metrics, from top to bottom in post-order.
450    pub plan_metrics: Vec<PlanMetrics>,
451    /// Per-region watermark for incremental-read checkpoint advancement.
452    ///
453    /// The watermark is the latest sequence (`seq`) this query round safely read
454    /// for each participating region. Flow uses it to decide where the next
455    /// incremental round can resume.
456    ///
457    /// - `Some(seq)`: the query proved it safely read up to `seq`; downstream
458    ///   may advance the checkpoint to this value.
459    /// - `None`: the region participated but the query could not prove a safe
460    ///   read upper-bound, so the checkpoint must not advance for this region.
461    ///
462    /// Omitted when empty for backward compatibility.
463    #[serde(default, skip_serializing_if = "Vec::is_empty")]
464    pub region_watermarks: Vec<RegionWatermarkEntry>,
465}
466
467#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
468pub struct RegionWatermarkEntry {
469    pub region_id: u64,
470    #[serde(default, skip_serializing_if = "Option::is_none")]
471    pub watermark: Option<u64>,
472}
473
474/// Determines if a metric name represents a time measurement that should be formatted.
475fn is_time_metric(metric_name: &str) -> bool {
476    metric_name.contains("elapsed") || metric_name.contains("time") || metric_name.contains("cost")
477}
478
479/// Determines if a metric name represents a bytes measurement that should be formatted.
480fn is_bytes_metric(metric_name: &str) -> bool {
481    metric_name.contains("bytes") || metric_name.contains("mem")
482}
483
484fn format_bytes_human_readable(bytes: usize) -> String {
485    format!("{}", ReadableSize(bytes as u64))
486}
487
488/// Only display `plan_metrics` with indent `  ` (2 spaces).
489impl Display for RecordBatchMetrics {
490    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
491        for metric in &self.plan_metrics {
492            write!(
493                f,
494                "{:indent$}{} metrics=[",
495                " ",
496                metric.plan.trim_end(),
497                indent = metric.level * 2,
498            )?;
499            for (label, value) in &metric.metrics {
500                if is_time_metric(label) {
501                    write!(
502                        f,
503                        "{}: {}, ",
504                        label,
505                        format_nanoseconds_human_readable(*value),
506                    )?;
507                } else if is_bytes_metric(label) {
508                    write!(f, "{}: {}, ", label, format_bytes_human_readable(*value),)?;
509                } else {
510                    write!(f, "{}: {}, ", label, value)?;
511                }
512            }
513            writeln!(f, "]")?;
514        }
515
516        Ok(())
517    }
518}
519
520#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
521pub struct PlanMetrics {
522    /// The plan name
523    pub plan: String,
524    /// The stable execution plan name.
525    #[serde(default)]
526    pub plan_name: String,
527    /// The level of the plan, starts from 0
528    pub level: usize,
529    /// An ordered key-value list of metrics.
530    /// Key is metric label and value is metric value.
531    pub metrics: Vec<(String, usize)>,
532}
533
534enum AsyncRecordBatchStreamAdapterState {
535    Uninit(FutureStream),
536    Ready(SendableRecordBatchStream),
537    Failed,
538}
539
540pub struct AsyncRecordBatchStreamAdapter {
541    schema: SchemaRef,
542    state: AsyncRecordBatchStreamAdapterState,
543}
544
545impl AsyncRecordBatchStreamAdapter {
546    pub fn new(schema: SchemaRef, stream: FutureStream) -> Self {
547        Self {
548            schema,
549            state: AsyncRecordBatchStreamAdapterState::Uninit(stream),
550        }
551    }
552}
553
554impl RecordBatchStream for AsyncRecordBatchStreamAdapter {
555    fn schema(&self) -> SchemaRef {
556        self.schema.clone()
557    }
558
559    fn output_ordering(&self) -> Option<&[OrderOption]> {
560        None
561    }
562
563    fn metrics(&self) -> Option<RecordBatchMetrics> {
564        None
565    }
566}
567
568impl Stream for AsyncRecordBatchStreamAdapter {
569    type Item = Result<RecordBatch>;
570
571    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
572        loop {
573            match &mut self.state {
574                AsyncRecordBatchStreamAdapterState::Uninit(stream_future) => {
575                    match ready!(Pin::new(stream_future).poll(cx)) {
576                        Ok(stream) => {
577                            self.state = AsyncRecordBatchStreamAdapterState::Ready(stream);
578                            continue;
579                        }
580                        Err(e) => {
581                            self.state = AsyncRecordBatchStreamAdapterState::Failed;
582                            return Poll::Ready(Some(Err(e)));
583                        }
584                    };
585                }
586                AsyncRecordBatchStreamAdapterState::Ready(stream) => {
587                    return Poll::Ready(ready!(Pin::new(stream).poll_next(cx)));
588                }
589                AsyncRecordBatchStreamAdapterState::Failed => return Poll::Ready(None),
590            }
591        }
592    }
593
594    // This is not supported for lazy stream.
595    #[inline]
596    fn size_hint(&self) -> (usize, Option<usize>) {
597        (0, None)
598    }
599}
600
601/// Custom cast function that handles Map -> Binary (JSON) conversion
602fn custom_cast(
603    array: &dyn Array,
604    target_type: &ArrowDataType,
605    extype: Option<ColumnExtType>,
606) -> std::result::Result<Arc<dyn Array>, ArrowError> {
607    if let ArrowDataType::Map(_, _) = array.data_type()
608        && let ArrowDataType::Binary = target_type
609    {
610        return convert_map_to_json_binary(array, extype);
611    }
612
613    cast(array, target_type)
614}
615
616/// Convert a Map array to a Binary array containing JSON data
617fn convert_map_to_json_binary(
618    array: &dyn Array,
619    extype: Option<ColumnExtType>,
620) -> std::result::Result<Arc<dyn Array>, ArrowError> {
621    use datatypes::arrow::array::{BinaryArray, MapArray};
622    use serde_json::Value;
623
624    let map_array = array
625        .as_any()
626        .downcast_ref::<MapArray>()
627        .ok_or_else(|| ArrowError::CastError("Failed to downcast to MapArray".to_string()))?;
628
629    let mut json_values = Vec::with_capacity(map_array.len());
630
631    for i in 0..map_array.len() {
632        if map_array.is_null(i) {
633            json_values.push(None);
634        } else {
635            // Extract the map entry at index i
636            let map_entry = map_array.value(i);
637            let key_value_array = map_entry
638                .as_any()
639                .downcast_ref::<datatypes::arrow::array::StructArray>()
640                .ok_or_else(|| {
641                    ArrowError::CastError("Failed to downcast to StructArray".to_string())
642                })?;
643
644            // Convert to JSON object
645            let mut json_obj = serde_json::Map::with_capacity(key_value_array.len());
646
647            for j in 0..key_value_array.len() {
648                if key_value_array.is_null(j) {
649                    continue;
650                }
651                let key_field = key_value_array.column(0);
652                let value_field = key_value_array.column(1);
653
654                if key_field.is_null(j) {
655                    continue;
656                }
657
658                let key = key_field
659                    .as_any()
660                    .downcast_ref::<datatypes::arrow::array::StringArray>()
661                    .ok_or_else(|| {
662                        ArrowError::CastError("Failed to downcast key to StringArray".to_string())
663                    })?
664                    .value(j);
665
666                let value = if value_field.is_null(j) {
667                    Value::Null
668                } else {
669                    let value_str = value_field
670                        .as_any()
671                        .downcast_ref::<datatypes::arrow::array::StringArray>()
672                        .ok_or_else(|| {
673                            ArrowError::CastError(
674                                "Failed to downcast value to StringArray".to_string(),
675                            )
676                        })?
677                        .value(j);
678                    Value::String(value_str.to_string())
679                };
680
681                json_obj.insert(key.to_string(), value);
682            }
683
684            let json_value = Value::Object(json_obj);
685            let json_bytes = match extype {
686                Some(ColumnExtType::Json) => {
687                    let json_string = match serde_json::to_string(&json_value) {
688                        Ok(s) => s,
689                        Err(e) => {
690                            return Err(ArrowError::CastError(format!(
691                                "Failed to serialize JSON: {}",
692                                e
693                            )));
694                        }
695                    };
696                    match jsonb::parse_value(json_string.as_bytes()) {
697                        Ok(jsonb_value) => jsonb_value.to_vec(),
698                        Err(e) => {
699                            return Err(ArrowError::CastError(format!(
700                                "Failed to serialize JSONB: {}",
701                                e
702                            )));
703                        }
704                    }
705                }
706                _ => match serde_json::to_vec(&json_value) {
707                    Ok(b) => b,
708                    Err(e) => {
709                        return Err(ArrowError::CastError(format!(
710                            "Failed to serialize JSON: {}",
711                            e
712                        )));
713                    }
714                },
715            };
716            json_values.push(Some(json_bytes));
717        }
718    }
719
720    let binary_array = BinaryArray::from_iter(json_values);
721    Ok(Arc::new(binary_array))
722}
723
724#[cfg(test)]
725mod test {
726    use common_error::ext::BoxedError;
727    use common_error::mock::MockError;
728    use common_error::status_code::StatusCode;
729    use datatypes::arrow::array::{ArrayRef, MapArray, StringArray, StructArray};
730    use datatypes::arrow::buffer::OffsetBuffer;
731    use datatypes::arrow::datatypes::Field;
732    use datatypes::prelude::ConcreteDataType;
733    use datatypes::schema::ColumnSchema;
734    use datatypes::vectors::Int32Vector;
735    use serde_json::json;
736    use snafu::IntoError;
737
738    use super::*;
739    use crate::RecordBatches;
740    use crate::error::Error;
741
742    #[tokio::test]
743    async fn test_async_recordbatch_stream_adaptor() {
744        struct MaybeErrorRecordBatchStream {
745            items: Vec<Result<RecordBatch>>,
746        }
747
748        impl RecordBatchStream for MaybeErrorRecordBatchStream {
749            fn schema(&self) -> SchemaRef {
750                unimplemented!()
751            }
752
753            fn output_ordering(&self) -> Option<&[OrderOption]> {
754                None
755            }
756
757            fn metrics(&self) -> Option<RecordBatchMetrics> {
758                None
759            }
760        }
761
762        impl Stream for MaybeErrorRecordBatchStream {
763            type Item = Result<RecordBatch>;
764
765            fn poll_next(
766                mut self: Pin<&mut Self>,
767                _: &mut Context<'_>,
768            ) -> Poll<Option<Self::Item>> {
769                if let Some(batch) = self.items.pop() {
770                    Poll::Ready(Some(Ok(batch?)))
771                } else {
772                    Poll::Ready(None)
773                }
774            }
775        }
776
777        fn new_future_stream(
778            maybe_recordbatches: Result<Vec<Result<RecordBatch>>>,
779        ) -> FutureStream {
780            Box::pin(async move {
781                maybe_recordbatches
782                    .map(|items| Box::pin(MaybeErrorRecordBatchStream { items }) as _)
783            })
784        }
785
786        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
787            "a",
788            ConcreteDataType::int32_datatype(),
789            false,
790        )]));
791        let batch1 = RecordBatch::new(
792            schema.clone(),
793            vec![Arc::new(Int32Vector::from_slice([1])) as _],
794        )
795        .unwrap();
796        let batch2 = RecordBatch::new(
797            schema.clone(),
798            vec![Arc::new(Int32Vector::from_slice([2])) as _],
799        )
800        .unwrap();
801
802        let success_stream = new_future_stream(Ok(vec![Ok(batch1.clone()), Ok(batch2.clone())]));
803        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), success_stream);
804        let collected = RecordBatches::try_collect(Box::pin(adapter)).await.unwrap();
805        assert_eq!(
806            collected,
807            RecordBatches::try_new(schema.clone(), vec![batch2.clone(), batch1.clone()]).unwrap()
808        );
809
810        let poll_err_stream = new_future_stream(Ok(vec![
811            Ok(batch1.clone()),
812            Err(error::ExternalSnafu
813                .into_error(BoxedError::new(MockError::new(StatusCode::Unknown)))),
814        ]));
815        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), poll_err_stream);
816        let err = RecordBatches::try_collect(Box::pin(adapter))
817            .await
818            .unwrap_err();
819        assert!(
820            matches!(err, Error::External { .. }),
821            "unexpected err {err}"
822        );
823
824        let failed_to_init_stream =
825            new_future_stream(Err(error::ExternalSnafu
826                .into_error(BoxedError::new(MockError::new(StatusCode::Internal)))));
827        let adapter = AsyncRecordBatchStreamAdapter::new(schema.clone(), failed_to_init_stream);
828        let err = RecordBatches::try_collect(Box::pin(adapter))
829            .await
830            .unwrap_err();
831        assert!(
832            matches!(err, Error::External { .. }),
833            "unexpected err {err}"
834        );
835    }
836
837    #[test]
838    fn test_convert_map_to_json_binary() {
839        let keys = StringArray::from(vec![Some("a"), Some("b"), Some("c"), Some("x")]);
840        let values = StringArray::from(vec![Some("1"), None, Some("3"), Some("42")]);
841        let key_field = Arc::new(Field::new("key", ArrowDataType::Utf8, false));
842        let value_field = Arc::new(Field::new("value", ArrowDataType::Utf8, true));
843        let struct_type = ArrowDataType::Struct(vec![key_field, value_field].into());
844
845        let entries_field = Arc::new(Field::new("entries", struct_type, false));
846
847        let struct_array = StructArray::from(vec![
848            (
849                Arc::new(Field::new("key", ArrowDataType::Utf8, false)),
850                Arc::new(keys) as ArrayRef,
851            ),
852            (
853                Arc::new(Field::new("value", ArrowDataType::Utf8, true)),
854                Arc::new(values) as ArrayRef,
855            ),
856        ]);
857
858        let offsets = OffsetBuffer::from_lengths([3, 0, 1]);
859        let nulls = datatypes::arrow::buffer::NullBuffer::from(vec![true, false, true]);
860
861        let map_array = MapArray::new(
862            entries_field,
863            offsets,
864            struct_array,
865            Some(nulls), // nulls
866            false,
867        );
868
869        let result = convert_map_to_json_binary(&map_array, None).unwrap();
870        let binary_array = result
871            .as_any()
872            .downcast_ref::<datatypes::arrow::array::BinaryArray>()
873            .unwrap();
874
875        let expected_jsons = [
876            Some(r#"{"a":"1","b":null,"c":"3"}"#),
877            None,
878            Some(r#"{"x":"42"}"#),
879        ];
880
881        for (i, _) in expected_jsons.iter().enumerate() {
882            if let Some(expected) = &expected_jsons[i] {
883                assert!(!binary_array.is_null(i));
884                let actual_bytes = binary_array.value(i);
885                let actual_str = std::str::from_utf8(actual_bytes).unwrap();
886                assert_eq!(actual_str, *expected);
887            } else {
888                assert!(binary_array.is_null(i));
889            }
890        }
891
892        let result_json =
893            convert_map_to_json_binary(&map_array, Some(ColumnExtType::Json)).unwrap();
894        let binary_array_json = result_json
895            .as_any()
896            .downcast_ref::<datatypes::arrow::array::BinaryArray>()
897            .unwrap();
898
899        for (i, _) in expected_jsons.iter().enumerate() {
900            if expected_jsons[i].is_some() {
901                assert!(!binary_array_json.is_null(i));
902                let actual_bytes = binary_array_json.value(i);
903                assert_ne!(actual_bytes, expected_jsons[i].unwrap().as_bytes());
904            } else {
905                assert!(binary_array_json.is_null(i));
906            }
907        }
908    }
909
910    #[test]
911    fn test_recordbatch_metrics_deserializes_without_region_watermarks() {
912        let metrics: RecordBatchMetrics = serde_json::from_value(json!({
913            "elapsed_compute": 12,
914            "memory_usage": 34,
915            "plan_metrics": []
916        }))
917        .unwrap();
918
919        assert!(metrics.region_watermarks.is_empty());
920        assert_eq!(metrics.elapsed_compute, 12);
921        assert_eq!(metrics.memory_usage, 34);
922    }
923
924    #[test]
925    fn test_plan_metrics_deserializes_without_plan_name() {
926        let metrics: RecordBatchMetrics = serde_json::from_value(json!({
927            "elapsed_compute": 12,
928            "memory_usage": 34,
929            "plan_metrics": [{
930                "plan": "SeqScan: region=1",
931                "level": 0,
932                "metrics": []
933            }]
934        }))
935        .unwrap();
936
937        assert_eq!(metrics.plan_metrics[0].plan_name, "");
938    }
939
940    #[test]
941    fn test_recordbatch_metrics_region_watermarks_serde_roundtrip() {
942        let metrics = RecordBatchMetrics {
943            region_watermarks: vec![
944                RegionWatermarkEntry {
945                    region_id: 1,
946                    watermark: Some(100),
947                },
948                RegionWatermarkEntry {
949                    region_id: 2,
950                    watermark: None,
951                },
952            ],
953            ..Default::default()
954        };
955
956        let value = serde_json::to_value(&metrics).unwrap();
957        assert_eq!(
958            value.get("region_watermarks").unwrap(),
959            &json!([
960                { "region_id": 1, "watermark": 100 },
961                { "region_id": 2 }
962            ])
963        );
964
965        let decoded: RecordBatchMetrics = serde_json::from_value(value).unwrap();
966        assert_eq!(decoded.region_watermarks, metrics.region_watermarks);
967    }
968
969    #[test]
970    fn test_recordbatch_metrics_skips_empty_region_watermarks_on_serialize() {
971        let value = serde_json::to_value(RecordBatchMetrics::default()).unwrap();
972        assert!(value.get("region_watermarks").is_none());
973    }
974}
common_recordbatch/adapter.rs

common_recordbatch/
adapter.rs
