connectorx/destinations/arrowstream/
mod.rs

1//! Destination implementation for Arrow and Polars.
2
3mod arrow_assoc;
4mod errors;
5mod funcs;
6pub mod typesystem;
7
8pub use self::errors::{ArrowDestinationError, Result};
9pub use self::typesystem::ArrowTypeSystem;
10use super::{Consume, Destination, DestinationPartition};
11use crate::constants::RECORD_BATCH_SIZE;
12use crate::data_order::DataOrder;
13use crate::errors::{ConnectorXOutError, OutResult};
14use crate::typesystem::{Realize, TypeAssoc, TypeSystem};
15use anyhow::anyhow;
16use arrow::{datatypes::Schema, record_batch::RecordBatch};
17use arrow_assoc::ArrowAssoc;
18use fehler::{throw, throws};
19use funcs::{FFinishBuilder, FNewBuilder, FNewField};
20use itertools::Itertools;
21use std::{
22    any::Any,
23    sync::{
24        mpsc::{channel, Receiver, Sender},
25        Arc,
26    },
27};
28
29type Builder = Box<dyn Any + Send>;
30type Builders = Vec<Builder>;
31
32/// A finished batch, or the error that stopped a partition.
33type Message = OutResult<RecordBatch>;
34
35pub struct ArrowDestination {
36    schema: Vec<ArrowTypeSystem>,
37    names: Vec<String>,
38    arrow_schema: Arc<Schema>,
39    batch_size: usize,
40    sender: Option<Sender<Message>>,
41    /// `None` once an error has been received: the partition writers then fail
42    /// on their next send instead of buffering batches nobody will read.
43    receiver: Option<Receiver<Message>>,
44}
45
46impl Default for ArrowDestination {
47    fn default() -> Self {
48        let (tx, rx) = channel();
49        ArrowDestination {
50            schema: vec![],
51            names: vec![],
52            arrow_schema: Arc::new(Schema::empty()),
53            batch_size: RECORD_BATCH_SIZE,
54            sender: Some(tx),
55            receiver: Some(rx),
56        }
57    }
58}
59
60impl ArrowDestination {
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    pub fn new_with_batch_size(batch_size: usize) -> Self {
66        let (tx, rx) = channel();
67        ArrowDestination {
68            schema: vec![],
69            names: vec![],
70            arrow_schema: Arc::new(Schema::empty()),
71            batch_size,
72            sender: Some(tx),
73            receiver: Some(rx),
74        }
75    }
76}
77
78impl Destination for ArrowDestination {
79    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::ColumnMajor, DataOrder::RowMajor];
80    type TypeSystem = ArrowTypeSystem;
81    type Partition<'a> = ArrowPartitionWriter;
82    type Error = ArrowDestinationError;
83
84    fn needs_count(&self) -> bool {
85        false
86    }
87
88    #[throws(ArrowDestinationError)]
89    fn allocate<S: AsRef<str>>(
90        &mut self,
91        _nrow: usize,
92        names: &[S],
93        schema: &[ArrowTypeSystem],
94        data_order: DataOrder,
95    ) {
96        // todo: support colmajor
97        if !matches!(data_order, DataOrder::RowMajor) {
98            throw!(crate::errors::ConnectorXError::UnsupportedDataOrder(
99                data_order
100            ))
101        }
102
103        // parse the metadata
104        self.schema = schema.to_vec();
105        self.names = names.iter().map(|n| n.as_ref().to_string()).collect();
106        let fields = self
107            .schema
108            .iter()
109            .zip_eq(&self.names)
110            .map(|(&dt, h)| Ok(Realize::<FNewField>::realize(dt)?(h.as_str())))
111            .collect::<Result<Vec<_>>>()?;
112        self.arrow_schema = Arc::new(Schema::new(fields));
113    }
114
115    #[throws(ArrowDestinationError)]
116    fn partition(&mut self, counts: usize) -> Vec<Self::Partition<'_>> {
117        let mut partitions = vec![];
118        let sender = self.sender.take().unwrap();
119        for _ in 0..counts {
120            partitions.push(ArrowPartitionWriter::new(
121                self.schema.clone(),
122                Arc::clone(&self.arrow_schema),
123                self.batch_size,
124                sender.clone(),
125            )?);
126        }
127        partitions
128        // self.sender should be freed
129    }
130
131    fn schema(&self) -> &[ArrowTypeSystem] {
132        self.schema.as_slice()
133    }
134}
135
136impl ArrowDestination {
137    #[throws(ConnectorXOutError)]
138    pub fn arrow(self) -> Vec<RecordBatch> {
139        if self.sender.is_some() {
140            // should not happen since it is dropped after partition
141            // but need to make sure here otherwise recv will be blocked forever
142            std::mem::drop(self.sender);
143        }
144        let mut data = vec![];
145        if let Some(receiver) = self.receiver {
146            while let Ok(rb) = receiver.recv() {
147                data.push(rb?);
148            }
149        }
150        data
151    }
152
153    /// Blocks until the next record batch is available.
154    ///
155    /// Returns `None` once every partition writer is done, or `Some(Err(_))`
156    /// as soon as a partition fails. After an error the channel is closed, so
157    /// later calls return `None` and the other partitions stop at their next
158    /// batch boundary.
159    pub fn record_batch(&mut self) -> Option<Message> {
160        let msg = self.receiver.as_ref()?.recv().ok()?;
161        if msg.is_err() {
162            self.receiver = None;
163        }
164        Some(msg)
165    }
166
167    pub fn empty_batch(&self) -> RecordBatch {
168        RecordBatch::new_empty(self.arrow_schema.clone())
169    }
170
171    pub fn arrow_schema(&self) -> Arc<Schema> {
172        self.arrow_schema.clone()
173    }
174
175    pub fn names(&self) -> &[String] {
176        self.names.as_slice()
177    }
178}
179
180pub struct ArrowPartitionWriter {
181    schema: Vec<ArrowTypeSystem>,
182    builders: Option<Builders>,
183    current_row: usize,
184    current_col: usize,
185    arrow_schema: Arc<Schema>,
186    batch_size: usize,
187    sender: Option<Sender<Message>>,
188}
189
190// unsafe impl Sync for ArrowPartitionWriter {}
191
192impl ArrowPartitionWriter {
193    #[throws(ArrowDestinationError)]
194    fn new(
195        schema: Vec<ArrowTypeSystem>,
196        arrow_schema: Arc<Schema>,
197        batch_size: usize,
198        sender: Sender<Message>,
199    ) -> Self {
200        let mut pw = ArrowPartitionWriter {
201            schema,
202            builders: None,
203            current_row: 0,
204            current_col: 0,
205            arrow_schema,
206            batch_size,
207            sender: Some(sender),
208        };
209        pw.allocate()?;
210        pw
211    }
212
213    #[throws(ArrowDestinationError)]
214    fn allocate(&mut self) {
215        let builders = self
216            .schema
217            .iter()
218            .map(|dt| Ok(Realize::<FNewBuilder>::realize(*dt)?(self.batch_size)))
219            .collect::<Result<Vec<_>>>()?;
220        self.builders.replace(builders);
221    }
222
223    #[throws(ArrowDestinationError)]
224    fn flush(&mut self) {
225        let builders = self
226            .builders
227            .take()
228            .unwrap_or_else(|| panic!("arrow builder is none when flush!"));
229        let columns = builders
230            .into_iter()
231            .zip(self.schema.iter())
232            .map(|(builder, &dt)| Realize::<FFinishBuilder>::realize(dt)?(builder))
233            .collect::<std::result::Result<Vec<_>, crate::errors::ConnectorXError>>()?;
234        let rb = RecordBatch::try_new(Arc::clone(&self.arrow_schema), columns)?;
235        if let Some(sender) = &self.sender {
236            // The consumer is gone (dropped, or already got an error from
237            // another partition): stop instead of parsing the rest for nothing.
238            sender
239                .send(Ok(rb))
240                .map_err(|_| anyhow!("record batch receiver dropped"))?;
241        }
242
243        self.current_row = 0;
244        self.current_col = 0;
245    }
246
247    /// Reports `err` to the consumer and closes this partition's side of the
248    /// stream. Rows buffered since the last flush are discarded.
249    ///
250    /// Call this while the source's query is still alive: some drivers drain
251    /// the remaining result set when the query is dropped (MySQL does), and
252    /// the consumer should not have to wait for that.
253    pub fn send_error(&mut self, err: ConnectorXOutError) {
254        self.builders = None;
255        if let Some(sender) = self.sender.take() {
256            // The consumer may already be gone; nothing else to do then.
257            let _ = sender.send(Err(err));
258        }
259    }
260}
261
262impl<'a> DestinationPartition<'a> for ArrowPartitionWriter {
263    type TypeSystem = ArrowTypeSystem;
264    type Error = ArrowDestinationError;
265
266    #[throws(ArrowDestinationError)]
267    fn finalize(&mut self) {
268        if self.builders.is_some() {
269            self.flush()?;
270        }
271        // need to release the sender so receiver knows when the stream is exhasted
272        std::mem::drop(self.sender.take());
273    }
274
275    #[throws(ArrowDestinationError)]
276    fn aquire_row(&mut self, _n: usize) -> usize {
277        self.current_row
278    }
279
280    fn ncols(&self) -> usize {
281        self.schema.len()
282    }
283}
284
285impl<'a, T> Consume<T> for ArrowPartitionWriter
286where
287    T: TypeAssoc<<Self as DestinationPartition<'a>>::TypeSystem> + ArrowAssoc + 'static,
288{
289    type Error = ArrowDestinationError;
290
291    #[throws(ArrowDestinationError)]
292    fn consume(&mut self, value: T) {
293        let col = self.current_col;
294        self.current_col = (self.current_col + 1) % self.ncols();
295        self.schema[col].check::<T>()?;
296
297        loop {
298            match &mut self.builders {
299                Some(builders) => {
300                    <T as ArrowAssoc>::append(
301                        builders[col]
302                            .downcast_mut::<T::Builder>()
303                            .ok_or_else(|| anyhow!("cannot cast arrow builder for append"))?,
304                        value,
305                    )?;
306                    break;
307                }
308                None => self.allocate()?, // allocate if builders are not initialized
309            }
310        }
311
312        // flush if exceed batch_size
313        if self.current_col == 0 {
314            self.current_row += 1;
315            if self.current_row >= self.batch_size {
316                self.flush()?;
317                self.allocate()?;
318            }
319        }
320    }
321}
connectorx/destinations/arrowstream/mod.rs

connectorx/destinations/arrowstream/
mod.rs
