monarch_record_batch/
lib.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Shared trait for Arrow RecordBatch buffer types.
10//!
11//! `RecordBatchBuffer` is the trait generated by
12//! `#[derive(RecordBatchRow)]` (from `monarch_record_batch_macros`).
13//! It is extracted here so the derive macro can be used from any
14//! Monarch crate, not only `monarch_distributed_telemetry`.
15
16use datafusion::arrow::error::ArrowError;
17use datafusion::arrow::record_batch::RecordBatch;
18// Re-export the derive macro so consumers only need to depend on this
19// crate.
20pub use monarch_record_batch_macros::RecordBatchRow;
21
22/// Trait for buffer types that can produce Arrow RecordBatches.
23///
24/// Auto-implemented by `#[derive(RecordBatchRow)]` from
25/// `monarch_record_batch_macros`.
26pub trait RecordBatchBuffer {
27    /// Number of rows currently buffered.
28    fn len(&self) -> usize;
29
30    /// Whether the buffer is empty.
31    fn is_empty(&self) -> bool {
32        self.len() == 0
33    }
34
35    /// Drain buffered rows into a [`RecordBatch`].
36    ///
37    /// After this call the buffer is empty (fields are consumed via
38    /// `std::mem::take`). The only failure mode is Arrow schema/batch
39    /// construction, so the error type is [`ArrowError`] rather than
40    /// a catch-all.
41    fn drain_to_record_batch(&mut self) -> Result<RecordBatch, ArrowError>;
42}
monarch_record_batch/lib.rs

monarch_record_batch/
lib.rs
