common_recordbatch/
lib.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
15pub mod adapter;
16pub mod cursor;
17pub mod error;
18pub mod ext;
19pub mod filter;
20mod memory;
21pub mod recordbatch;
22pub mod util;
23
24use std::fmt::{self, Write};
25use std::pin::Pin;
26use std::sync::Arc;
27
28use adapter::RecordBatchMetrics;
29use arc_swap::ArcSwapOption;
30use common_base::readable_size::ReadableSize;
31use common_memory_manager::{MemoryManager, MemoryMetrics, PermitGranularity};
32use common_telemetry::tracing::Span;
33pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
34use datatypes::arrow::array::{Array, ArrayRef, AsArray, StringBuilder};
35use datatypes::arrow::compute::SortOptions;
36use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field};
37use datatypes::arrow::error::ArrowError;
38pub use datatypes::arrow::record_batch::RecordBatch as DfRecordBatch;
39use datatypes::arrow::util::display::{
40    ArrayFormatter, ArrayFormatterFactory, DisplayIndex, FormatOptions, FormatResult,
41};
42use datatypes::arrow::util::pretty::{
43    pretty_format_batches_with_options, pretty_format_batches_with_schema,
44};
45use datatypes::extension::json::is_any_json_extension_type;
46use datatypes::prelude::{ConcreteDataType, DataType, VectorRef};
47use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
48use datatypes::types::{JsonFormat, StructField, StructType, jsonb_to_string};
49use error::Result;
50use futures::task::{Context, Poll};
51use futures::{Stream, TryStreamExt};
52pub use recordbatch::RecordBatch;
53use snafu::{ResultExt, ensure};
54
55use crate::error::{ArrowComputeSnafu, NewDfRecordBatchSnafu};
56pub use crate::memory::QueryMemoryReservation;
57
58pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
59    fn name(&self) -> &str {
60        "RecordBatchStream"
61    }
62
63    fn schema(&self) -> SchemaRef;
64
65    fn output_ordering(&self) -> Option<&[OrderOption]>;
66
67    fn metrics(&self) -> Option<RecordBatchMetrics>;
68}
69
70pub type SendableRecordBatchStream = Pin<Box<dyn RecordBatchStream + Send>>;
71
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct OrderOption {
74    pub name: String,
75    pub options: SortOptions,
76}
77
78/// A wrapper that maps a [RecordBatchStream] to a new [RecordBatchStream] by applying a function to each [RecordBatch].
79///
80/// The mapper function is applied to each [RecordBatch] in the stream.
81/// The schema of the new [RecordBatchStream] is the same as the schema of the inner [RecordBatchStream] after applying the schema mapper function.
82/// The output ordering of the new [RecordBatchStream] is the same as the output ordering of the inner [RecordBatchStream].
83/// The metrics of the new [RecordBatchStream] is the same as the metrics of the inner [RecordBatchStream] if it is not `None`.
84pub struct SendableRecordBatchMapper {
85    inner: SendableRecordBatchStream,
86    /// The mapper function is applied to each [RecordBatch] in the stream.
87    /// The original schema and the mapped schema are passed to the mapper function.
88    mapper: fn(RecordBatch, &SchemaRef, &SchemaRef) -> Result<RecordBatch>,
89    /// The schema of the new [RecordBatchStream] is the same as the schema of the inner [RecordBatchStream] after applying the schema mapper function.
90    schema: SchemaRef,
91    /// Whether the mapper function is applied to each [RecordBatch] in the stream.
92    apply_mapper: bool,
93}
94
95/// Maps the json type to string in the batch.
96///
97/// The json type is mapped to string by converting the json value to string.
98/// The batch is updated to have the same number of columns as the original batch,
99/// but with the json type mapped to string.
100pub fn map_json_type_to_string(
101    batch: RecordBatch,
102    original_schema: &SchemaRef,
103    mapped_schema: &SchemaRef,
104) -> Result<RecordBatch> {
105    let mut vectors = Vec::with_capacity(original_schema.column_schemas().len());
106    for (vector, schema) in batch.columns().iter().zip(original_schema.column_schemas()) {
107        if let ConcreteDataType::Json(j) = &schema.data_type {
108            if matches!(&j.format, JsonFormat::Jsonb) {
109                let mut string_vector_builder = StringBuilder::new();
110                let binary_vector = vector.as_binary::<i32>();
111                for value in binary_vector.iter() {
112                    let Some(value) = value else {
113                        string_vector_builder.append_null();
114                        continue;
115                    };
116                    let string_value =
117                        jsonb_to_string(value).with_context(|_| error::CastVectorSnafu {
118                            from_type: schema.data_type.clone(),
119                            to_type: ConcreteDataType::string_datatype(),
120                        })?;
121                    string_vector_builder.append_value(string_value);
122                }
123
124                let string_vector = string_vector_builder.finish();
125                vectors.push(Arc::new(string_vector) as ArrayRef);
126            } else {
127                vectors.push(vector.clone());
128            }
129        } else {
130            vectors.push(vector.clone());
131        }
132    }
133
134    let record_batch = datatypes::arrow::record_batch::RecordBatch::try_new(
135        mapped_schema.arrow_schema().clone(),
136        vectors,
137    )
138    .context(NewDfRecordBatchSnafu)?;
139    Ok(RecordBatch::from_df_record_batch(
140        mapped_schema.clone(),
141        record_batch,
142    ))
143}
144
145/// Maps the json type to string in the schema.
146///
147/// The json type is mapped to string by converting the json value to string.
148/// The schema is updated to have the same number of columns as the original schema,
149/// but with the json type mapped to string.
150///
151/// Returns the new schema and whether the schema needs to be mapped to string.
152pub fn map_json_type_to_string_schema(schema: SchemaRef) -> (SchemaRef, bool) {
153    let mut new_columns = Vec::with_capacity(schema.column_schemas().len());
154    let mut apply_mapper = false;
155    for column in schema.column_schemas() {
156        if matches!(column.data_type, ConcreteDataType::Json(_)) {
157            new_columns.push(ColumnSchema::new(
158                column.name.clone(),
159                ConcreteDataType::string_datatype(),
160                column.is_nullable(),
161            ));
162            apply_mapper = true;
163        } else {
164            new_columns.push(column.clone());
165        }
166    }
167    (Arc::new(Schema::new(new_columns)), apply_mapper)
168}
169
170/// Replaces dictionary types with their value types, including dictionaries nested in lists and
171/// structs.
172pub fn map_dictionary_to_values_data_type(data_type: &ConcreteDataType) -> ConcreteDataType {
173    match data_type {
174        ConcreteDataType::Dictionary(dictionary) => {
175            map_dictionary_to_values_data_type(dictionary.value_type())
176        }
177        ConcreteDataType::List(list) => ConcreteDataType::list_datatype(Arc::new(
178            map_dictionary_to_values_data_type(list.item_type()),
179        )),
180        ConcreteDataType::Struct(struct_type) => {
181            let fields = struct_type
182                .fields()
183                .iter()
184                .map(|field| {
185                    StructField::new(
186                        field.name(),
187                        map_dictionary_to_values_data_type(field.data_type()),
188                        field.is_nullable(),
189                    )
190                })
191                .collect();
192            ConcreteDataType::struct_datatype(StructType::new(Arc::new(fields)))
193        }
194        _ => data_type.clone(),
195    }
196}
197
198/// Maps dictionary columns in a schema to their value types.
199pub fn map_dictionary_to_values_schema(schema: SchemaRef) -> (SchemaRef, bool) {
200    let mut apply_mapper = false;
201    let columns: Vec<_> = schema
202        .column_schemas()
203        .iter()
204        .map(|column| {
205            let data_type = map_dictionary_to_values_data_type(&column.data_type);
206            apply_mapper |= data_type != column.data_type;
207            let mut column = column.clone();
208            column.data_type = data_type;
209            column
210        })
211        .collect();
212
213    if !apply_mapper {
214        return (schema, false);
215    }
216
217    // Query projections may contain duplicate column names, which SchemaBuilder rejects. Preserve
218    // the existing Arrow fields and only replace their data types so field and schema metadata are
219    // retained as well.
220    let fields = schema
221        .arrow_schema()
222        .fields()
223        .iter()
224        .zip(schema.column_schemas())
225        .zip(&columns)
226        .map(|((field, original), column)| {
227            if original.data_type == column.data_type {
228                field.clone()
229            } else {
230                Arc::new(
231                    field
232                        .as_ref()
233                        .clone()
234                        .with_data_type(column.data_type.as_arrow_type()),
235                )
236            }
237        })
238        .collect::<Vec<_>>();
239    let arrow_schema =
240        datatypes::arrow::datatypes::Schema::new(fields).with_metadata(schema.metadata().clone());
241    (
242        Arc::new(Schema::try_from(Arc::new(arrow_schema)).unwrap()),
243        true,
244    )
245}
246
247/// Expands dictionary arrays to their value arrays according to `mapped_schema`.
248pub fn map_dictionary_to_values(
249    batch: RecordBatch,
250    original_schema: &SchemaRef,
251    mapped_schema: &SchemaRef,
252) -> Result<RecordBatch> {
253    let arrays = batch
254        .columns()
255        .iter()
256        .zip(original_schema.column_schemas())
257        .zip(mapped_schema.column_schemas())
258        .map(|((array, original), mapped)| {
259            if original.data_type == mapped.data_type {
260                Ok(array.clone())
261            } else {
262                datatypes::arrow::compute::cast(array, &mapped.data_type.as_arrow_type())
263                    .context(ArrowComputeSnafu)
264            }
265        })
266        .collect::<Result<Vec<_>>>()?;
267
268    let record_batch = DfRecordBatch::try_new(mapped_schema.arrow_schema().clone(), arrays)
269        .context(NewDfRecordBatchSnafu)?;
270    Ok(RecordBatch::from_df_record_batch(
271        mapped_schema.clone(),
272        record_batch,
273    ))
274}
275
276impl SendableRecordBatchMapper {
277    /// Creates a new [SendableRecordBatchMapper] with the given inner [RecordBatchStream], mapper function, and schema mapper function.
278    pub fn new(
279        inner: SendableRecordBatchStream,
280        mapper: fn(RecordBatch, &SchemaRef, &SchemaRef) -> Result<RecordBatch>,
281        schema_mapper: fn(SchemaRef) -> (SchemaRef, bool),
282    ) -> Self {
283        let (mapped_schema, apply_mapper) = schema_mapper(inner.schema());
284        Self {
285            inner,
286            mapper,
287            schema: mapped_schema,
288            apply_mapper,
289        }
290    }
291}
292
293impl RecordBatchStream for SendableRecordBatchMapper {
294    fn name(&self) -> &str {
295        "SendableRecordBatchMapper"
296    }
297
298    fn schema(&self) -> SchemaRef {
299        self.schema.clone()
300    }
301
302    fn output_ordering(&self) -> Option<&[OrderOption]> {
303        self.inner.output_ordering()
304    }
305
306    fn metrics(&self) -> Option<RecordBatchMetrics> {
307        self.inner.metrics()
308    }
309}
310
311impl Stream for SendableRecordBatchMapper {
312    type Item = Result<RecordBatch>;
313
314    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
315        if self.apply_mapper {
316            Pin::new(&mut self.inner).poll_next(cx).map(|opt| {
317                opt.map(|result| {
318                    result
319                        .and_then(|batch| (self.mapper)(batch, &self.inner.schema(), &self.schema))
320                })
321            })
322        } else {
323            Pin::new(&mut self.inner).poll_next(cx)
324        }
325    }
326}
327
328/// EmptyRecordBatchStream can be used to create a RecordBatchStream
329/// that will produce no results
330pub struct EmptyRecordBatchStream {
331    /// Schema wrapped by Arc
332    schema: SchemaRef,
333}
334
335impl EmptyRecordBatchStream {
336    /// Create an empty RecordBatchStream
337    pub fn new(schema: SchemaRef) -> Self {
338        Self { schema }
339    }
340}
341
342impl RecordBatchStream for EmptyRecordBatchStream {
343    fn schema(&self) -> SchemaRef {
344        self.schema.clone()
345    }
346
347    fn output_ordering(&self) -> Option<&[OrderOption]> {
348        None
349    }
350
351    fn metrics(&self) -> Option<RecordBatchMetrics> {
352        None
353    }
354}
355
356impl Stream for EmptyRecordBatchStream {
357    type Item = Result<RecordBatch>;
358
359    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
360        Poll::Ready(None)
361    }
362}
363
364#[derive(Debug, PartialEq)]
365pub struct RecordBatches {
366    schema: SchemaRef,
367    batches: Vec<RecordBatch>,
368}
369
370impl RecordBatches {
371    pub fn try_from_columns<I: IntoIterator<Item = VectorRef>>(
372        schema: SchemaRef,
373        columns: I,
374    ) -> Result<Self> {
375        let batches = vec![RecordBatch::new(schema.clone(), columns)?];
376        Ok(Self { schema, batches })
377    }
378
379    pub async fn try_collect(stream: SendableRecordBatchStream) -> Result<Self> {
380        let schema = stream.schema();
381        let batches = stream.try_collect::<Vec<_>>().await?;
382        Ok(Self { schema, batches })
383    }
384
385    #[inline]
386    pub fn empty() -> Self {
387        Self {
388            schema: Arc::new(Schema::new(vec![])),
389            batches: vec![],
390        }
391    }
392
393    pub fn iter(&self) -> impl Iterator<Item = &RecordBatch> {
394        self.batches.iter()
395    }
396
397    pub fn pretty_print(&self) -> Result<String> {
398        let df_batches = &self
399            .iter()
400            .map(|x| x.df_record_batch().clone())
401            .collect::<Vec<_>>();
402        let result: String = if df_batches.is_empty() {
403            pretty_format_batches_with_schema(self.schema.arrow_schema().clone(), df_batches)
404                .context(error::FormatSnafu)?
405                .to_string()
406        } else {
407            let options =
408                FormatOptions::default().with_formatter_factory(Some(&BinaryFormatterFactory));
409            pretty_format_batches_with_options(df_batches, &options)
410                .context(error::FormatSnafu)?
411                .to_string()
412        };
413
414        Ok(result)
415    }
416
417    pub fn try_new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Result<Self> {
418        for batch in &batches {
419            ensure!(
420                batch.schema == schema,
421                error::CreateRecordBatchesSnafu {
422                    reason: format!(
423                        "expect RecordBatch schema equals {:?}, actual: {:?}",
424                        schema, batch.schema
425                    )
426                }
427            )
428        }
429        Ok(Self { schema, batches })
430    }
431
432    pub fn schema(&self) -> SchemaRef {
433        self.schema.clone()
434    }
435
436    pub fn take(self) -> Vec<RecordBatch> {
437        self.batches
438    }
439
440    pub fn as_stream(&self) -> SendableRecordBatchStream {
441        Box::pin(SimpleRecordBatchStream {
442            inner: RecordBatches {
443                schema: self.schema(),
444                batches: self.batches.clone(),
445            },
446            index: 0,
447        })
448    }
449}
450
451#[derive(Debug)]
452struct BinaryFormatterFactory;
453
454impl ArrayFormatterFactory for BinaryFormatterFactory {
455    fn create_array_formatter<'a>(
456        &self,
457        array: &'a dyn Array,
458        options: &FormatOptions<'a>,
459        field: Option<&'a Field>,
460    ) -> std::result::Result<Option<ArrayFormatter<'a>>, ArrowError> {
461        if !array.data_type().is_binary() {
462            return Ok(None);
463        }
464
465        Ok(Some(ArrayFormatter::new(
466            Box::new(BinaryFormatter {
467                array,
468                is_json: field.is_some_and(is_any_json_extension_type),
469                default: ArrayFormatter::try_new(array, options)?,
470                null: options.null(),
471            }),
472            options.safe(),
473        )))
474    }
475}
476
477struct BinaryFormatter<'a> {
478    array: &'a dyn Array,
479    is_json: bool,
480    default: ArrayFormatter<'a>,
481    null: &'a str,
482}
483
484impl DisplayIndex for BinaryFormatter<'_> {
485    fn write(&self, idx: usize, f: &mut dyn Write) -> FormatResult {
486        if !self.is_json {
487            self.default.value(idx).write(f)?;
488            return Ok(());
489        }
490
491        if self.array.is_null(idx) {
492            write!(f, "{}", self.null)?;
493        } else {
494            let bytes = match self.array.data_type() {
495                ArrowDataType::Binary => self.array.as_binary::<i32>().value(idx),
496                ArrowDataType::LargeBinary => self.array.as_binary::<i64>().value(idx),
497                ArrowDataType::BinaryView => self.array.as_binary_view().value(idx),
498                _ => return Ok(self.default.value(idx).write(f)?),
499            };
500            let value =
501                jsonb_to_string(bytes).map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
502            write!(f, "{value}")?;
503        }
504        Ok(())
505    }
506}
507
508impl IntoIterator for RecordBatches {
509    type Item = RecordBatch;
510    type IntoIter = std::vec::IntoIter<Self::Item>;
511
512    fn into_iter(self) -> Self::IntoIter {
513        self.batches.into_iter()
514    }
515}
516
517pub struct SimpleRecordBatchStream {
518    inner: RecordBatches,
519    index: usize,
520}
521
522impl RecordBatchStream for SimpleRecordBatchStream {
523    fn schema(&self) -> SchemaRef {
524        self.inner.schema()
525    }
526
527    fn output_ordering(&self) -> Option<&[OrderOption]> {
528        None
529    }
530
531    fn metrics(&self) -> Option<RecordBatchMetrics> {
532        None
533    }
534}
535
536impl Stream for SimpleRecordBatchStream {
537    type Item = Result<RecordBatch>;
538
539    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
540        Poll::Ready(if self.index < self.inner.batches.len() {
541            let batch = self.inner.batches[self.index].clone();
542            self.index += 1;
543            Some(Ok(batch))
544        } else {
545            None
546        })
547    }
548}
549
550/// Adapt a [Stream] of [RecordBatch] to a [RecordBatchStream].
551pub struct RecordBatchStreamWrapper<S> {
552    pub schema: SchemaRef,
553    pub stream: S,
554    pub output_ordering: Option<Vec<OrderOption>>,
555    pub metrics: Arc<ArcSwapOption<RecordBatchMetrics>>,
556    pub span: Span,
557}
558
559impl<S> RecordBatchStreamWrapper<S> {
560    /// Creates a [RecordBatchStreamWrapper] without output ordering requirement.
561    pub fn new(schema: SchemaRef, stream: S) -> RecordBatchStreamWrapper<S> {
562        RecordBatchStreamWrapper {
563            schema,
564            stream,
565            output_ordering: None,
566            metrics: Default::default(),
567            span: Span::current(),
568        }
569    }
570}
571
572impl<S: Stream<Item = Result<RecordBatch>> + Unpin> RecordBatchStream
573    for RecordBatchStreamWrapper<S>
574{
575    fn name(&self) -> &str {
576        "RecordBatchStreamWrapper"
577    }
578
579    fn schema(&self) -> SchemaRef {
580        self.schema.clone()
581    }
582
583    fn output_ordering(&self) -> Option<&[OrderOption]> {
584        self.output_ordering.as_deref()
585    }
586
587    fn metrics(&self) -> Option<RecordBatchMetrics> {
588        self.metrics.load().as_ref().map(|s| s.as_ref().clone())
589    }
590}
591
592impl<S: Stream<Item = Result<RecordBatch>> + Unpin> Stream for RecordBatchStreamWrapper<S> {
593    type Item = Result<RecordBatch>;
594
595    fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
596        let _entered = self.span.clone().entered();
597        Pin::new(&mut self.stream).poll_next(ctx)
598    }
599}
600
601/// Memory tracker for RecordBatch streams and reader reservations. Clone to share
602/// the same limit across queries. Limited accounting rounds up to KB; unlimited
603/// accounting is byte-exact. Exhaustion fails immediately. DataFusion operators use
604/// the separate query runtime pool.
605#[derive(Clone)]
606pub struct QueryMemoryTracker {
607    manager: MemoryManager<CallbackMemoryMetrics>,
608    metrics: CallbackMemoryMetrics,
609}
610
611impl fmt::Debug for QueryMemoryTracker {
612    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
613        f.debug_struct("QueryMemoryTracker")
614            .field("current", &self.current())
615            .field("limit", &self.limit())
616            .field("on_update", &self.metrics.has_on_update())
617            .field("on_exhausted", &self.metrics.has_on_exhausted())
618            .field("on_rejected", &self.metrics.has_on_rejected())
619            .finish()
620    }
621}
622
623impl QueryMemoryTracker {
624    /// Create a builder for a query memory tracker.
625    pub fn builder(limit: usize) -> QueryMemoryTrackerBuilder {
626        QueryMemoryTrackerBuilder {
627            limit,
628            on_update: None,
629            on_exhausted: None,
630            on_reject: None,
631        }
632    }
633
634    fn new_stream_tracker(&self) -> StreamMemoryTracker {
635        StreamMemoryTracker {
636            tracker: self.clone(),
637            guard: QueryMemoryReservation {
638                // A zero-byte acquisition always succeeds; this manager is never closed.
639                inner: self.manager.try_acquire(0).unwrap(),
640            },
641            tracked_bytes: 0,
642        }
643    }
644    /// Get the current memory usage in bytes.
645    pub fn current(&self) -> usize {
646        self.manager.used_bytes() as usize
647    }
648
649    /// Returns the memory limit in bytes, or zero when unlimited.
650    pub fn limit(&self) -> usize {
651        self.manager.limit_bytes() as usize
652    }
653
654    /// Reserves bytes until the returned guard is dropped, failing immediately if
655    /// the shared budget cannot accommodate the reservation.
656    pub fn reserve(&self, bytes: usize) -> Result<QueryMemoryReservation> {
657        let mut tracker = self.new_stream_tracker();
658        if let Err(error) = tracker.try_track(bytes) {
659            tracker.inc_rejected();
660            return Err(error);
661        }
662        Ok(tracker.guard)
663    }
664
665    fn reject_error(
666        &self,
667        current: usize,
668        additional: usize,
669        stream_tracked: usize,
670    ) -> error::Error {
671        let limit = self.limit();
672        let msg = format!(
673            "{} requested, {} used globally ({}%), {} used by this stream, hard limit: {}",
674            ReadableSize(additional as u64),
675            ReadableSize(current as u64),
676            current.saturating_mul(100).checked_div(limit).unwrap_or(0),
677            ReadableSize(stream_tracked as u64),
678            ReadableSize(limit as u64)
679        );
680        error::ExceedMemoryLimitSnafu { msg }.build()
681    }
682
683    fn inc_rejected(&self) {
684        self.metrics.inc_rejected();
685    }
686}
687
688/// Builder for constructing a [`QueryMemoryTracker`] with optional callbacks.
689pub struct QueryMemoryTrackerBuilder {
690    limit: usize,
691    on_update: Option<UpdateCallback>,
692    on_exhausted: Option<UnitCallback>,
693    on_reject: Option<RejectCallback>,
694}
695
696impl QueryMemoryTrackerBuilder {
697    /// Set a callback to be called whenever the usage changes successfully.
698    /// The callback receives the new total usage in bytes.
699    ///
700    /// # Note
701    /// The callback is called after both successful `track()` and stream drop.
702    /// Usage is byte-exact in unlimited mode; otherwise it is 1KB-aligned.
703    pub fn on_update<F>(mut self, on_update: F) -> Self
704    where
705        F: Fn(usize) + Send + Sync + 'static,
706    {
707        self.on_update = Some(Arc::new(on_update));
708        self
709    }
710
711    /// Set a callback to be called when memory is unavailable for immediate acquisition.
712    ///
713    /// # Note
714    /// This is called when the non-blocking allocation fast path fails.
715    /// It is never called when `limit == 0` (unlimited mode).
716    pub fn on_exhausted<F>(mut self, on_exhausted: F) -> Self
717    where
718        F: Fn() + Send + Sync + 'static,
719    {
720        self.on_exhausted = Some(Arc::new(on_exhausted));
721        self
722    }
723
724    /// Set a callback to be called when the request ultimately fails due to memory pressure.
725    pub fn on_reject<F>(mut self, on_reject: F) -> Self
726    where
727        F: Fn() + Send + Sync + 'static,
728    {
729        self.on_reject = Some(Arc::new(on_reject));
730        self
731    }
732
733    /// Build a [`QueryMemoryTracker`] from this builder.
734    pub fn build(self) -> QueryMemoryTracker {
735        let metrics = CallbackMemoryMetrics::new(self.on_update, self.on_exhausted, self.on_reject);
736        let manager = MemoryManager::with_granularity(
737            self.limit as u64,
738            PermitGranularity::Kilobyte,
739            metrics.clone(),
740        );
741
742        QueryMemoryTracker { manager, metrics }
743    }
744}
745
746struct StreamMemoryTracker {
747    tracker: QueryMemoryTracker,
748    guard: QueryMemoryReservation,
749    tracked_bytes: usize,
750}
751
752impl StreamMemoryTracker {
753    fn inc_rejected(&self) {
754        self.tracker.inc_rejected();
755    }
756
757    fn try_track(&mut self, additional: usize) -> Result<()> {
758        if self.guard.try_acquire_additional(additional as u64) {
759            self.tracked_bytes = self.tracked_bytes.saturating_add(additional);
760            Ok(())
761        } else {
762            Err(self.reject_error(additional))
763        }
764    }
765
766    fn reject_error(&self, additional: usize) -> error::Error {
767        let current = self.tracker.current();
768        self.tracker
769            .reject_error(current, additional, self.tracked_bytes)
770    }
771}
772
773#[derive(Clone)]
774struct CallbackMemoryMetrics {
775    inner: Arc<CallbackMemoryMetricsInner>,
776}
777
778type UpdateCallback = Arc<dyn Fn(usize) + Send + Sync>;
779type UnitCallback = Arc<dyn Fn() + Send + Sync>;
780type RejectCallback = UnitCallback;
781
782struct CallbackMemoryMetricsInner {
783    on_update: Option<UpdateCallback>,
784    on_exhausted: Option<UnitCallback>,
785    on_reject: Option<RejectCallback>,
786}
787
788impl CallbackMemoryMetrics {
789    fn new(
790        on_update: Option<UpdateCallback>,
791        on_exhausted: Option<UnitCallback>,
792        on_reject: Option<RejectCallback>,
793    ) -> Self {
794        Self {
795            inner: Arc::new(CallbackMemoryMetricsInner {
796                on_update,
797                on_exhausted,
798                on_reject,
799            }),
800        }
801    }
802
803    fn has_on_update(&self) -> bool {
804        self.inner.on_update.is_some()
805    }
806
807    fn has_on_exhausted(&self) -> bool {
808        self.inner.on_exhausted.is_some()
809    }
810
811    fn has_on_rejected(&self) -> bool {
812        self.inner.on_reject.is_some()
813    }
814
815    fn inc_rejected(&self) {
816        if let Some(callback) = &self.inner.on_reject {
817            callback();
818        }
819    }
820}
821
822impl MemoryMetrics for CallbackMemoryMetrics {
823    fn set_limit(&self, _: i64) {}
824
825    fn set_in_use(&self, bytes: i64) {
826        if let Some(callback) = &self.inner.on_update {
827            callback(bytes.max(0) as usize);
828        }
829    }
830
831    fn inc_exhausted(&self, _: &str) {
832        if let Some(callback) = &self.inner.on_exhausted {
833            callback();
834        }
835    }
836}
837
838/// A wrapper stream that tracks memory usage of RecordBatches.
839pub struct MemoryTrackedStream {
840    inner: SendableRecordBatchStream,
841    tracker: StreamMemoryTracker,
842}
843
844impl MemoryTrackedStream {
845    pub fn new(inner: SendableRecordBatchStream, tracker: QueryMemoryTracker) -> Self {
846        Self {
847            inner,
848            tracker: tracker.new_stream_tracker(),
849        }
850    }
851
852    fn track_batch(&mut self, batch: RecordBatch) -> Result<RecordBatch> {
853        let additional = batch.logical_slice_memory_size();
854        if let Err(error) = self.tracker.try_track(additional) {
855            self.tracker.inc_rejected();
856            return Err(error);
857        }
858        Ok(batch)
859    }
860}
861
862impl Stream for MemoryTrackedStream {
863    type Item = Result<RecordBatch>;
864
865    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
866        match Pin::new(&mut self.inner).poll_next(cx) {
867            Poll::Ready(Some(Ok(batch))) => Poll::Ready(Some(self.track_batch(batch))),
868            Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))),
869            Poll::Ready(None) => Poll::Ready(None),
870            Poll::Pending => Poll::Pending,
871        }
872    }
873
874    fn size_hint(&self) -> (usize, Option<usize>) {
875        self.inner.size_hint()
876    }
877}
878
879impl RecordBatchStream for MemoryTrackedStream {
880    fn schema(&self) -> SchemaRef {
881        self.inner.schema()
882    }
883
884    fn output_ordering(&self) -> Option<&[OrderOption]> {
885        self.inner.output_ordering()
886    }
887
888    fn metrics(&self) -> Option<RecordBatchMetrics> {
889        self.inner.metrics()
890    }
891}
892
893#[cfg(test)]
894mod tests {
895    use std::sync::Arc;
896    use std::sync::atomic::{AtomicUsize, Ordering};
897
898    use common_memory_manager::PermitGranularity;
899    use datatypes::arrow::array::{
900        DictionaryArray, Int32Array, ListArray, StringArray, UInt32Array,
901    };
902    use datatypes::arrow::buffer::OffsetBuffer;
903    use datatypes::arrow::datatypes::{
904        DataType as ArrowDataType, Field, Int32Type, Schema as ArrowSchema,
905    };
906    use datatypes::prelude::{ConcreteDataType, VectorRef};
907    use datatypes::schema::{ColumnSchema, Schema};
908    use datatypes::vectors::{BooleanVector, Int32Vector, StringVector};
909    use futures::StreamExt;
910
911    use super::*;
912
913    fn large_string_batch(bytes: usize) -> RecordBatch {
914        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
915            "payload",
916            ConcreteDataType::string_datatype(),
917            false,
918        )]));
919        let payload = "x".repeat(bytes);
920        let vector: VectorRef = Arc::new(StringVector::from(vec![payload]));
921        RecordBatch::new(schema, vec![vector]).unwrap()
922    }
923
924    fn aligned_tracked_bytes(bytes: usize) -> usize {
925        PermitGranularity::Kilobyte
926            .permits_to_bytes(PermitGranularity::Kilobyte.bytes_to_permits(bytes as u64))
927            as usize
928    }
929
930    #[test]
931    fn test_reader_reservation_rejects_and_releases_capacity() {
932        let tracker = QueryMemoryTracker::builder(2048).build();
933        let first = tracker.reserve(1025).unwrap();
934        assert_eq!(2048, tracker.current());
935        assert!(tracker.reserve(1).is_err());
936        assert_eq!(2048, tracker.current());
937        drop(first);
938        assert_eq!(0, tracker.current());
939
940        assert!(tracker.reserve(2049).is_err());
941        assert_eq!(0, tracker.current());
942        let guard = tracker.reserve(2048).unwrap();
943        assert_eq!(2048, tracker.current());
944        drop(guard);
945        assert_eq!(0, tracker.current());
946    }
947
948    #[tokio::test]
949    async fn test_reader_and_batch_reservations_share_scan_budget() {
950        let batch = large_string_batch(512);
951        let batch_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size());
952        let tracker = QueryMemoryTracker::builder(batch_bytes + 1024).build();
953        let reader = tracker.reserve(1).unwrap();
954        let mut stream = MemoryTrackedStream::new(
955            RecordBatches::try_new(batch.schema.clone(), vec![batch.clone(), batch])
956                .unwrap()
957                .as_stream(),
958            tracker.clone(),
959        );
960        stream.next().await.unwrap().unwrap();
961        assert_eq!(tracker.limit(), tracker.current());
962        assert!(tracker.reserve(1).is_err());
963        assert!(stream.next().await.unwrap().is_err());
964        drop(reader);
965        assert_eq!(batch_bytes, tracker.current());
966        drop(stream);
967        assert_eq!(0, tracker.current());
968        let unlimited = QueryMemoryTracker::builder(0).build();
969        let guard = unlimited.reserve(17).unwrap();
970        assert_eq!(17, unlimited.current());
971        drop(guard);
972        assert_eq!(0, unlimited.current());
973    }
974
975    #[tokio::test]
976    async fn test_memory_tracked_stream_charges_logical_slice_size() {
977        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
978            "payload",
979            ConcreteDataType::string_datatype(),
980            false,
981        )]));
982        let payloads: Vec<_> = (0..1024)
983            .map(|value| format!("payload-{value:04}"))
984            .collect();
985        let batch = RecordBatch::new(schema, vec![Arc::new(StringVector::from(payloads)) as _])
986            .unwrap()
987            .slice(512, 1)
988            .unwrap();
989        let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size());
990        assert!(expected_bytes < aligned_tracked_bytes(batch.buffer_memory_size()));
991        let tracker = QueryMemoryTracker::builder(MB).build();
992        let mut stream = MemoryTrackedStream::new(
993            RecordBatches::try_new(batch.schema.clone(), vec![batch])
994                .unwrap()
995                .as_stream(),
996            tracker.clone(),
997        );
998
999        stream.next().await.unwrap().unwrap();
1000
1001        assert_eq!(tracker.current(), expected_bytes);
1002    }
1003
1004    #[test]
1005    fn test_recordbatches_try_from_columns() {
1006        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
1007            "a",
1008            ConcreteDataType::int32_datatype(),
1009            false,
1010        )]));
1011        let result = RecordBatches::try_from_columns(
1012            schema.clone(),
1013            vec![Arc::new(StringVector::from(vec!["hello", "world"])) as _],
1014        );
1015        assert!(result.is_err());
1016
1017        let v: VectorRef = Arc::new(Int32Vector::from_slice([1, 2]));
1018        let expected = vec![RecordBatch::new(schema.clone(), vec![v.clone()]).unwrap()];
1019        let r = RecordBatches::try_from_columns(schema, vec![v]).unwrap();
1020        assert_eq!(r.take(), expected);
1021    }
1022
1023    #[tokio::test]
1024    async fn test_recordbatches_pretty_print_empty_batches_preserves_schema() {
1025        let schema = Arc::new(Schema::new(vec![
1026            ColumnSchema::new("unit", ConcreteDataType::string_datatype(), false),
1027            ColumnSchema::new(
1028                "ts",
1029                ConcreteDataType::timestamp_millisecond_datatype(),
1030                false,
1031            ),
1032            ColumnSchema::new(
1033                "lhs.degrees(val) + rhs.radians(val)",
1034                ConcreteDataType::float64_datatype(),
1035                false,
1036            ),
1037        ]));
1038        let batches =
1039            RecordBatches::try_collect(Box::pin(EmptyRecordBatchStream::new(schema.clone())))
1040                .await
1041                .unwrap();
1042
1043        assert_eq!(schema, batches.schema());
1044        let expected = "\
1045+------+----+-------------------------------------+
1046| unit | ts | lhs.degrees(val) + rhs.radians(val) |
1047+------+----+-------------------------------------+
1048+------+----+-------------------------------------+";
1049        assert_eq!(expected, batches.pretty_print().unwrap());
1050    }
1051
1052    #[test]
1053    fn test_recordbatches_try_new() {
1054        let column_a = ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false);
1055        let column_b = ColumnSchema::new("b", ConcreteDataType::string_datatype(), false);
1056        let column_c = ColumnSchema::new("c", ConcreteDataType::boolean_datatype(), false);
1057
1058        let va: VectorRef = Arc::new(Int32Vector::from_slice([1, 2]));
1059        let vb: VectorRef = Arc::new(StringVector::from(vec!["hello", "world"]));
1060        let vc: VectorRef = Arc::new(BooleanVector::from(vec![true, false]));
1061
1062        let schema1 = Arc::new(Schema::new(vec![column_a.clone(), column_b]));
1063        let batch1 = RecordBatch::new(schema1.clone(), vec![va.clone(), vb]).unwrap();
1064
1065        let schema2 = Arc::new(Schema::new(vec![column_a, column_c]));
1066        let batch2 = RecordBatch::new(schema2.clone(), vec![va, vc]).unwrap();
1067
1068        let result = RecordBatches::try_new(schema1.clone(), vec![batch1.clone(), batch2]);
1069        assert!(result.is_err());
1070        assert_eq!(
1071            result.unwrap_err().to_string(),
1072            format!(
1073                "Failed to create RecordBatches, reason: expect RecordBatch schema equals {schema1:?}, actual: {schema2:?}",
1074            )
1075        );
1076
1077        let batches = RecordBatches::try_new(schema1.clone(), vec![batch1.clone()]).unwrap();
1078        let expected = "\
1079+---+-------+
1080| a | b     |
1081+---+-------+
1082| 1 | hello |
1083| 2 | world |
1084+---+-------+";
1085        assert_eq!(batches.pretty_print().unwrap(), expected);
1086
1087        assert_eq!(schema1, batches.schema());
1088        assert_eq!(vec![batch1], batches.take());
1089    }
1090
1091    #[test]
1092    fn test_map_dictionary_to_values_recursively() {
1093        let string_dictionary = ConcreteDataType::dictionary_datatype(
1094            ConcreteDataType::int32_datatype(),
1095            ConcreteDataType::string_datatype(),
1096        );
1097        let list_dictionary =
1098            ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::dictionary_datatype(
1099                ConcreteDataType::uint32_datatype(),
1100                ConcreteDataType::string_datatype(),
1101            )));
1102        let schema = Arc::new(Schema::new(vec![
1103            ColumnSchema::new("host", string_dictionary, true),
1104            ColumnSchema::new("tags", list_dictionary, true),
1105        ]));
1106
1107        let host = DictionaryArray::<Int32Type>::new(
1108            Int32Array::from(vec![0, 1]),
1109            Arc::new(StringArray::from(vec![Some("host-a"), None])),
1110        );
1111        let ArrowDataType::List(item_field) = schema.arrow_schema().field(1).data_type().clone()
1112        else {
1113            unreachable!()
1114        };
1115        let tag_values = DictionaryArray::new(
1116            UInt32Array::from_iter_values([0, 1, 0]),
1117            Arc::new(StringArray::from(vec![Some("a"), None])),
1118        );
1119        let tags = ListArray::new(
1120            item_field,
1121            OffsetBuffer::from_lengths([2, 1]),
1122            Arc::new(tag_values),
1123            None,
1124        );
1125        let batch = DfRecordBatch::try_new(
1126            schema.arrow_schema().clone(),
1127            vec![Arc::new(host), Arc::new(tags)],
1128        )
1129        .unwrap();
1130        let batch = RecordBatch::from_df_record_batch(schema.clone(), batch);
1131
1132        let (mapped_schema, apply_mapper) = map_dictionary_to_values_schema(schema.clone());
1133        assert!(apply_mapper);
1134        assert_eq!(
1135            &ArrowDataType::Utf8,
1136            mapped_schema.arrow_schema().field(0).data_type()
1137        );
1138        assert_eq!(
1139            &ArrowDataType::List(Arc::new(
1140                datatypes::arrow::datatypes::Field::new_list_field(ArrowDataType::Utf8, true,)
1141            )),
1142            mapped_schema.arrow_schema().field(1).data_type()
1143        );
1144
1145        let mapped = map_dictionary_to_values(batch, &schema, &mapped_schema).unwrap();
1146        let host = mapped
1147            .column(0)
1148            .as_any()
1149            .downcast_ref::<StringArray>()
1150            .unwrap();
1151        assert_eq!(vec![Some("host-a"), None], host.iter().collect::<Vec<_>>());
1152        let tags = mapped
1153            .column(1)
1154            .as_any()
1155            .downcast_ref::<ListArray>()
1156            .unwrap();
1157        let tag_values = tags
1158            .values()
1159            .as_any()
1160            .downcast_ref::<StringArray>()
1161            .unwrap();
1162        assert_eq!(
1163            vec![Some("a"), None, Some("a")],
1164            tag_values.iter().collect::<Vec<_>>()
1165        );
1166    }
1167
1168    #[test]
1169    fn test_map_dictionary_to_values_preserves_timestamp_timezone() {
1170        let dictionary = DictionaryArray::<Int32Type>::new(
1171            Int32Array::from(vec![Some(0), None]),
1172            Arc::new(StringArray::from(vec!["a"])),
1173        );
1174        let timestamps = datatypes::arrow::array::TimestampSecondArray::from(vec![Some(42), None])
1175            .with_timezone("Europe/Paris");
1176        let arrow_schema = Arc::new(ArrowSchema::new(vec![
1177            Field::new("label", dictionary.data_type().clone(), true),
1178            Field::new("instant", timestamps.data_type().clone(), true),
1179        ]));
1180        let schema = Arc::new(Schema::try_from(arrow_schema.clone()).unwrap());
1181        let batch = DfRecordBatch::try_new(
1182            arrow_schema.clone(),
1183            vec![Arc::new(dictionary), Arc::new(timestamps)],
1184        )
1185        .unwrap();
1186        let batch = RecordBatch::from_df_record_batch(schema.clone(), batch);
1187
1188        let (mapped_schema, apply_mapper) = map_dictionary_to_values_schema(schema.clone());
1189        assert!(apply_mapper);
1190        assert_eq!(arrow_schema.field(1), mapped_schema.arrow_schema().field(1));
1191        let mapped = map_dictionary_to_values(batch.clone(), &schema, &mapped_schema).unwrap();
1192        assert_eq!(&ArrowDataType::Utf8, mapped.column(0).data_type());
1193        assert_eq!(
1194            &StringArray::from(vec![Some("a"), None]),
1195            mapped
1196                .column(0)
1197                .as_any()
1198                .downcast_ref::<StringArray>()
1199                .unwrap()
1200        );
1201        assert!(Arc::ptr_eq(batch.column(1), mapped.column(1)));
1202    }
1203
1204    #[test]
1205    fn test_map_dictionary_to_values_schema_with_duplicate_columns() {
1206        let dictionary_type = ArrowDataType::Dictionary(
1207            Box::new(ArrowDataType::Int32),
1208            Box::new(ArrowDataType::Utf8),
1209        );
1210        let arrow_schema = Arc::new(ArrowSchema::new(vec![
1211            Field::new("area", dictionary_type.clone(), true),
1212            Field::new("area", dictionary_type, true),
1213        ]));
1214        let schema = Arc::new(Schema::try_from(arrow_schema).unwrap());
1215
1216        let (mapped_schema, apply_mapper) = map_dictionary_to_values_schema(schema);
1217        assert!(apply_mapper);
1218        assert_eq!(2, mapped_schema.num_columns());
1219        for field in mapped_schema.arrow_schema().fields() {
1220            assert_eq!("area", field.name());
1221            assert_eq!(&ArrowDataType::Utf8, field.data_type());
1222        }
1223    }
1224
1225    #[tokio::test]
1226    async fn test_simple_recordbatch_stream() {
1227        let column_a = ColumnSchema::new("a", ConcreteDataType::int32_datatype(), false);
1228        let column_b = ColumnSchema::new("b", ConcreteDataType::string_datatype(), false);
1229        let schema = Arc::new(Schema::new(vec![column_a, column_b]));
1230
1231        let va1: VectorRef = Arc::new(Int32Vector::from_slice([1, 2]));
1232        let vb1: VectorRef = Arc::new(StringVector::from(vec!["a", "b"]));
1233        let batch1 = RecordBatch::new(schema.clone(), vec![va1, vb1]).unwrap();
1234
1235        let va2: VectorRef = Arc::new(Int32Vector::from_slice([3, 4, 5]));
1236        let vb2: VectorRef = Arc::new(StringVector::from(vec!["c", "d", "e"]));
1237        let batch2 = RecordBatch::new(schema.clone(), vec![va2, vb2]).unwrap();
1238
1239        let recordbatches =
1240            RecordBatches::try_new(schema.clone(), vec![batch1.clone(), batch2.clone()]).unwrap();
1241        let stream = recordbatches.as_stream();
1242        let collected = util::collect(stream).await.unwrap();
1243        assert_eq!(collected.len(), 2);
1244        assert_eq!(collected[0], batch1);
1245        assert_eq!(collected[1], batch2);
1246    }
1247
1248    const MB: usize = 1024 * 1024;
1249
1250    #[test]
1251    fn test_query_memory_tracker_basic() {
1252        let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
1253
1254        let mut stream1 = tracker.new_stream_tracker();
1255        assert!(stream1.try_track(5 * MB).is_ok());
1256        assert_eq!(tracker.current(), 5 * MB);
1257
1258        let mut stream2 = tracker.new_stream_tracker();
1259        assert!(stream2.try_track(4 * MB).is_ok());
1260        assert_eq!(tracker.current(), 9 * MB);
1261
1262        drop(stream1);
1263        drop(stream2);
1264        assert_eq!(tracker.current(), 0);
1265    }
1266
1267    #[test]
1268    fn test_query_memory_tracker_shared_global_limit() {
1269        let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
1270        let mut stream1 = tracker.new_stream_tracker();
1271        let mut stream2 = tracker.new_stream_tracker();
1272
1273        assert!(stream1.try_track(3 * MB).is_ok());
1274        assert_eq!(tracker.current(), 3 * MB);
1275        assert!(stream2.try_track(6 * MB).is_ok());
1276        assert_eq!(tracker.current(), 9 * MB);
1277
1278        let err = stream2.try_track(2 * MB).unwrap_err();
1279        let err_msg = err.to_string();
1280        assert!(err_msg.contains("6.0MiB used by this stream"));
1281        assert!(err_msg.contains("9.0MiB used globally (90%)"));
1282        assert!(err_msg.contains("hard limit: 10.0MiB"));
1283        assert_eq!(tracker.current(), 9 * MB);
1284
1285        drop(stream1);
1286        assert_eq!(tracker.current(), 6 * MB);
1287        drop(stream2);
1288        assert_eq!(tracker.current(), 0);
1289    }
1290
1291    #[test]
1292    fn test_query_memory_tracker_hard_limit() {
1293        let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
1294        let mut stream = tracker.new_stream_tracker();
1295
1296        assert!(stream.try_track(9 * MB).is_ok());
1297        assert_eq!(tracker.current(), 9 * MB);
1298
1299        assert!(stream.try_track(2 * MB).is_err());
1300        assert_eq!(tracker.current(), 9 * MB);
1301
1302        assert!(stream.try_track(MB).is_ok());
1303        assert_eq!(tracker.current(), 10 * MB);
1304
1305        assert!(stream.try_track(MB).is_err());
1306        assert_eq!(tracker.current(), 10 * MB);
1307
1308        drop(stream);
1309        assert_eq!(tracker.current(), 0);
1310    }
1311
1312    #[test]
1313    fn test_query_memory_tracker_unlimited() {
1314        let tracker = Arc::new(QueryMemoryTracker::builder(0).build());
1315        let mut stream = tracker.new_stream_tracker();
1316
1317        assert!(stream.try_track(10 * MB).is_ok());
1318        assert_eq!(tracker.current(), 10 * MB);
1319        drop(stream);
1320        assert_eq!(tracker.current(), 0);
1321    }
1322
1323    #[test]
1324    fn test_query_memory_tracker_rounds_to_kilobytes() {
1325        let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
1326        let mut stream = tracker.new_stream_tracker();
1327
1328        assert!(stream.try_track(1_537).is_ok());
1329        assert_eq!(tracker.current(), 2 * 1024);
1330
1331        drop(stream);
1332        assert_eq!(tracker.current(), 0);
1333    }
1334
1335    #[tokio::test]
1336    async fn test_memory_tracked_stream_rejects_immediately() {
1337        let exhausted = Arc::new(AtomicUsize::new(0));
1338        let rejected = Arc::new(AtomicUsize::new(0));
1339        let exhausted_counter = exhausted.clone();
1340        let rejected_counter = rejected.clone();
1341        let tracker = QueryMemoryTracker::builder(MB)
1342            .on_exhausted(move || {
1343                exhausted_counter.fetch_add(1, Ordering::Relaxed);
1344            })
1345            .on_reject(move || {
1346                rejected_counter.fetch_add(1, Ordering::Relaxed);
1347            })
1348            .build();
1349        let batch = large_string_batch(700 * 1024);
1350
1351        let mut stream1 = MemoryTrackedStream::new(
1352            RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()])
1353                .unwrap()
1354                .as_stream(),
1355            tracker.clone(),
1356        );
1357        let first = stream1.next().await.unwrap().unwrap();
1358        assert_eq!(first.num_rows(), 1);
1359
1360        let mut stream2 = MemoryTrackedStream::new(
1361            RecordBatches::try_new(batch.schema.clone(), vec![batch])
1362                .unwrap()
1363                .as_stream(),
1364            tracker,
1365        );
1366        assert!(matches!(
1367            futures::poll!(stream2.next()),
1368            Poll::Ready(Some(Err(_)))
1369        ));
1370        assert_eq!(exhausted.load(Ordering::Relaxed), 1);
1371        assert_eq!(rejected.load(Ordering::Relaxed), 1);
1372    }
1373}
common_recordbatch/lib.rs

common_recordbatch/
lib.rs
