connectorx/destinations/arrow/
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::typesystem::{Realize, TypeAssoc, TypeSystem};
14use anyhow::anyhow;
15use arrow::{datatypes::Schema, record_batch::RecordBatch};
16use arrow_assoc::ArrowAssoc;
17use fehler::{throw, throws};
18use funcs::{FFinishBuilder, FNewBuilder, FNewField};
19use itertools::Itertools;
20use std::{
21    any::Any,
22    sync::{Arc, Mutex},
23};
24
25#[cfg(feature = "dst_polars")]
26use {
27    arrow::ffi::to_ffi,
28    polars::prelude::{concat, Column, DataFrame, IntoLazy, PlSmallStr, Series, UnionArgs},
29    polars_arrow::ffi::{import_array_from_c, import_field_from_c},
30    std::mem::transmute,
31};
32
33type Builder = Box<dyn Any + Send>;
34type Builders = Vec<Builder>;
35
36pub struct ArrowDestination {
37    schema: Vec<ArrowTypeSystem>,
38    names: Vec<String>,
39    data: Arc<Mutex<Vec<RecordBatch>>>,
40    arrow_schema: Arc<Schema>,
41    batch_size: usize,
42}
43
44impl Default for ArrowDestination {
45    fn default() -> Self {
46        ArrowDestination {
47            schema: vec![],
48            names: vec![],
49            data: Arc::new(Mutex::new(vec![])),
50            arrow_schema: Arc::new(Schema::empty()),
51            batch_size: RECORD_BATCH_SIZE,
52        }
53    }
54}
55
56impl ArrowDestination {
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    pub fn new_with_batch_size(batch_size: usize) -> Self {
62        ArrowDestination {
63            schema: vec![],
64            names: vec![],
65            data: Arc::new(Mutex::new(vec![])),
66            arrow_schema: Arc::new(Schema::empty()),
67            batch_size,
68        }
69    }
70}
71
72impl Destination for ArrowDestination {
73    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::ColumnMajor, DataOrder::RowMajor];
74    type TypeSystem = ArrowTypeSystem;
75    type Partition<'a> = ArrowPartitionWriter;
76    type Error = ArrowDestinationError;
77
78    fn needs_count(&self) -> bool {
79        false
80    }
81
82    #[throws(ArrowDestinationError)]
83    fn allocate<S: AsRef<str>>(
84        &mut self,
85        _nrow: usize,
86        names: &[S],
87        schema: &[ArrowTypeSystem],
88        data_order: DataOrder,
89    ) {
90        // todo: support colmajor
91        if !matches!(data_order, DataOrder::RowMajor) {
92            throw!(crate::errors::ConnectorXError::UnsupportedDataOrder(
93                data_order
94            ))
95        }
96
97        // parse the metadata
98        self.schema = schema.to_vec();
99        self.names = names.iter().map(|n| n.as_ref().to_string()).collect();
100        let fields = self
101            .schema
102            .iter()
103            .zip_eq(&self.names)
104            .map(|(&dt, h)| Ok(Realize::<FNewField>::realize(dt)?(h.as_str())))
105            .collect::<Result<Vec<_>>>()?;
106        self.arrow_schema = Arc::new(Schema::new(fields));
107    }
108
109    #[throws(ArrowDestinationError)]
110    fn partition(&mut self, counts: usize) -> Vec<Self::Partition<'_>> {
111        let mut partitions = vec![];
112        for _ in 0..counts {
113            partitions.push(ArrowPartitionWriter::new(
114                self.schema.clone(),
115                Arc::clone(&self.data),
116                Arc::clone(&self.arrow_schema),
117                self.batch_size,
118            )?);
119        }
120        partitions
121    }
122
123    fn schema(&self) -> &[ArrowTypeSystem] {
124        self.schema.as_slice()
125    }
126}
127
128impl ArrowDestination {
129    #[throws(ArrowDestinationError)]
130    pub fn arrow(self) -> Vec<RecordBatch> {
131        let lock = Arc::try_unwrap(self.data).map_err(|_| anyhow!("Partitions are not freed"))?;
132        lock.into_inner()
133            .map_err(|e| anyhow!("mutex poisoned {}", e))?
134    }
135
136    #[cfg(feature = "dst_polars")]
137    #[throws(ArrowDestinationError)]
138    pub fn polars(self) -> DataFrame {
139        // Convert to arrow first
140        let rbs = self.arrow()?;
141
142        // Ready LazyFrame vector for the chunks
143        let mut lf_vec = vec![];
144
145        for chunk in rbs.into_iter() {
146            // Column vector
147            let mut columns = Vec::with_capacity(chunk.num_columns());
148
149            // Arrow stores data by columns, therefore need to be Zero-copied by column
150            for (i, col) in chunk.columns().iter().enumerate() {
151                // Convert to ArrayData (arrow-rs)
152                let array = col.to_data();
153
154                // Convert to ffi with arrow-rs
155                let (out_array, out_schema) = to_ffi(&array).unwrap();
156
157                // Import field from ffi with polars
158                let field = unsafe {
159                    import_field_from_c(transmute::<
160                        &arrow::ffi::FFI_ArrowSchema,
161                        &polars_arrow::ffi::ArrowSchema,
162                    >(&out_schema))
163                }?;
164
165                // Import data from ffi with polars
166                let data = unsafe {
167                    import_array_from_c(
168                        transmute::<arrow::ffi::FFI_ArrowArray, polars_arrow::ffi::ArrowArray>(
169                            out_array,
170                        ),
171                        field.dtype().clone(),
172                    )
173                }?;
174
175                // Create Polars series from arrow column
176                columns.push(Column::from(Series::from_arrow(
177                    PlSmallStr::from(chunk.schema().field(i).name()),
178                    data,
179                )?));
180            }
181
182            // Create DataFrame from the columns
183            lf_vec.push(DataFrame::new(chunk.num_rows(), columns)?.lazy());
184        }
185
186        // Concat the chunks
187        let union_args = UnionArgs::default();
188        concat(lf_vec, union_args)?.collect()?
189    }
190
191    #[throws(ArrowDestinationError)]
192    pub fn record_batch(&mut self) -> Option<RecordBatch> {
193        let mut guard = self
194            .data
195            .lock()
196            .map_err(|e| anyhow!("mutex poisoned {}", e))?;
197        (*guard).pop()
198    }
199
200    pub fn empty_batch(&self) -> RecordBatch {
201        RecordBatch::new_empty(self.arrow_schema.clone())
202    }
203
204    pub fn arrow_schema(&self) -> Arc<Schema> {
205        self.arrow_schema.clone()
206    }
207
208    pub fn names(&self) -> &[String] {
209        self.names.as_slice()
210    }
211}
212
213pub struct ArrowPartitionWriter {
214    schema: Vec<ArrowTypeSystem>,
215    builders: Option<Builders>,
216    current_row: usize,
217    current_col: usize,
218    data: Arc<Mutex<Vec<RecordBatch>>>,
219    arrow_schema: Arc<Schema>,
220    batch_size: usize,
221}
222
223// unsafe impl Sync for ArrowPartitionWriter {}
224
225impl ArrowPartitionWriter {
226    #[throws(ArrowDestinationError)]
227    fn new(
228        schema: Vec<ArrowTypeSystem>,
229        data: Arc<Mutex<Vec<RecordBatch>>>,
230        arrow_schema: Arc<Schema>,
231        batch_size: usize,
232    ) -> Self {
233        let mut pw = ArrowPartitionWriter {
234            schema,
235            builders: None,
236            current_row: 0,
237            current_col: 0,
238            data,
239            arrow_schema,
240            batch_size,
241        };
242        pw.allocate()?;
243        pw
244    }
245
246    #[throws(ArrowDestinationError)]
247    fn allocate(&mut self) {
248        let builders = self
249            .schema
250            .iter()
251            .map(|dt| Ok(Realize::<FNewBuilder>::realize(*dt)?(self.batch_size)))
252            .collect::<Result<Vec<_>>>()?;
253        self.builders.replace(builders);
254    }
255
256    #[throws(ArrowDestinationError)]
257    fn flush(&mut self) {
258        let builders = self
259            .builders
260            .take()
261            .unwrap_or_else(|| panic!("arrow builder is none when flush!"));
262        let columns = builders
263            .into_iter()
264            .zip(self.schema.iter())
265            .map(|(builder, &dt)| Realize::<FFinishBuilder>::realize(dt)?(builder))
266            .collect::<std::result::Result<Vec<_>, crate::errors::ConnectorXError>>()?;
267        let rb = RecordBatch::try_new(Arc::clone(&self.arrow_schema), columns)?;
268        {
269            let mut guard = self
270                .data
271                .lock()
272                .map_err(|e| anyhow!("mutex poisoned {}", e))?;
273            let inner_data = &mut *guard;
274            inner_data.push(rb);
275        }
276
277        self.current_row = 0;
278        self.current_col = 0;
279    }
280}
281
282impl<'a> DestinationPartition<'a> for ArrowPartitionWriter {
283    type TypeSystem = ArrowTypeSystem;
284    type Error = ArrowDestinationError;
285
286    #[throws(ArrowDestinationError)]
287    fn finalize(&mut self) {
288        if self.builders.is_some() {
289            self.flush()?;
290        }
291    }
292
293    #[throws(ArrowDestinationError)]
294    fn aquire_row(&mut self, _n: usize) -> usize {
295        self.current_row
296    }
297
298    fn ncols(&self) -> usize {
299        self.schema.len()
300    }
301}
302
303impl<'a, T> Consume<T> for ArrowPartitionWriter
304where
305    T: TypeAssoc<<Self as DestinationPartition<'a>>::TypeSystem> + ArrowAssoc + 'static,
306{
307    type Error = ArrowDestinationError;
308
309    #[throws(ArrowDestinationError)]
310    fn consume(&mut self, value: T) {
311        let col = self.current_col;
312        self.current_col = (self.current_col + 1) % self.ncols();
313        self.schema[col].check::<T>()?;
314
315        loop {
316            match &mut self.builders {
317                Some(builders) => {
318                    <T as ArrowAssoc>::append(
319                        builders[col]
320                            .downcast_mut::<T::Builder>()
321                            .ok_or_else(|| anyhow!("cannot cast arrow builder for append"))?,
322                        value,
323                    )?;
324                    break;
325                }
326                None => self.allocate()?, // allocate if builders are not initialized
327            }
328        }
329
330        // flush if exceed batch_size
331        if self.current_col == 0 {
332            self.current_row += 1;
333            if self.current_row >= self.batch_size {
334                self.flush()?;
335                self.allocate()?;
336            }
337        }
338    }
339}
connectorx/destinations/arrow/mod.rs

connectorx/destinations/arrow/
mod.rs
